from __future__ import annotations
from typing import List, Iterable, Sequence
import re

from ..models.segments import Segment, SegmentSummary
from ..config import settings

# Rough token estimator (since we don't pull in tiktoken by default)
_DEF_CHARS_PER_TOKEN = 4

def _approx_tokens(text: str) -> int:
    return max(1, len(text) // _DEF_CHARS_PER_TOKEN)

# Heuristic: join sentences until target token window reached, respecting hard max.
# Provide resilience if text already large.

def _simple_sentence_tokenize(text: str) -> List[str]:
    # Basic sentence splitter; avoids heavy dependencies.
    # Split on punctuation followed by space/newline; keep punctuation.
    raw = re.split(r'(?<=[.!?])\s+', text.strip())
    return [r.strip() for r in raw if r.strip()]


def improved_segments(text: str, target_tokens: int = 180, hard_max_tokens: int = 260) -> List[Segment]:
    sentences = _simple_sentence_tokenize(text)
    segments: List[Segment] = []
    buf: List[str] = []
    buf_tokens = 0
    idx = 0
    for sent in sentences:
        stoks = _approx_tokens(sent)
        # If single sentence itself exceeds hard_max, flush previous and store alone.
        if stoks >= hard_max_tokens:
            if buf:
                seg_text = " ".join(buf).strip()
                segments.append(Segment(index=len(segments), text=seg_text, token_count=buf_tokens))
                buf, buf_tokens = [], 0
            segments.append(Segment(index=len(segments), text=sent.strip(), token_count=stoks, kind_hint="oversize"))
            continue
        # If adding sentence would exceed hard_max outright, flush buffer.
        if buf and buf_tokens + stoks > hard_max_tokens:
            seg_text = " ".join(buf).strip()
            segments.append(Segment(index=len(segments), text=seg_text, token_count=buf_tokens))
            buf, buf_tokens = [], 0
        buf.append(sent)
        buf_tokens += stoks
        # If we reached target window, flush.
        if buf_tokens >= target_tokens:
            seg_text = " ".join(buf).strip()
            segments.append(Segment(index=len(segments), text=seg_text, token_count=buf_tokens))
            buf, buf_tokens = [], 0
    if buf:
        seg_text = " ".join(buf).strip()
        segments.append(Segment(index=len(segments), text=seg_text, token_count=buf_tokens))
    return segments

_DEF_FALLBACK_SUMMARY_LEN = 42

def _fallback_micro_summary(text: str) -> str:
    # First sentence or trimmed leading fragment.
    first_period = text.find('.')
    candidate = text if first_period == -1 else text[: first_period + 1]
    candidate = re.sub('\s+', ' ', candidate).strip()
    if len(candidate) > 160:
        candidate = candidate[:157] + '...'
    return candidate


def batch_summarize_segments(segments: Sequence[Segment]) -> List[SegmentSummary]:
    cfg = settings
    # If LLM not enabled, fallback summarization.
    if not cfg.llm_enabled:
        return [
            SegmentSummary(
                segment_index=s.index,
                micro_summary=_fallback_micro_summary(s.text),
                key_points=[],
                retrieval_qas=[],
            )
            for s in segments
        ]
    # Lazy import to avoid mandatory dependency if user disables LLM.
    try:
        from ..llm import get_provider  # type: ignore
    except Exception:  # pragma: no cover - overly defensive
        return [
            SegmentSummary(
                segment_index=s.index,
                micro_summary=_fallback_micro_summary(s.text),
                key_points=[],
                retrieval_qas=[],
            )
            for s in segments
        ]

    provider = get_provider(cfg.llm_provider)
    if provider is None:
        return [
            SegmentSummary(
                segment_index=s.index,
                micro_summary=_fallback_micro_summary(s.text),
                key_points=[],
                retrieval_qas=[],
            )
            for s in segments
        ]
    # Construct a compact prompt with numbered segments asking for JSON lines.
    seg_payload = [f"[{s.index}] {s.text[:500]}" for s in segments]
    prompt = (
        "You are a concise learning designer. For each numbered segment, produce a JSON object with "
        "fields: segment_index, micro_summary (<=160 chars), key_points (<=4 short bullet strings), retrieval_qas (<=2 objects with q and a). "
        "Return one JSON object per line. If unsure, keep arrays empty.\nSegments:\n" + "\n".join(seg_payload)
    )

    try:
        # Provider expects an LLMRequest dataclass.
        from ..llm.base import LLMRequest  # local import
        raw_resp = provider.complete(LLMRequest(prompt=prompt))
        raw = getattr(raw_resp, 'text', str(raw_resp))
    except Exception:
        # Fallback gracefully
        return [
            SegmentSummary(
                segment_index=s.index,
                micro_summary=_fallback_micro_summary(s.text),
                key_points=[],
                retrieval_qas=[],
            )
            for s in segments
        ]

    # Parse line-delimited JSON objects; degrade gracefully on errors.
    import json
    summaries: List[SegmentSummary] = []
    for line in raw.splitlines():
        line = line.strip()
        if not line or not line.startswith('{'):
            continue
        try:
            obj = json.loads(line)
            if not isinstance(obj, dict):
                continue
            idx = obj.get('segment_index')
            if isinstance(idx, int):
                summaries.append(
                    SegmentSummary(
                        segment_index=idx,
                        micro_summary=obj.get('micro_summary') or _fallback_micro_summary(segments[idx].text),
                        key_points=[kp for kp in obj.get('key_points', []) if isinstance(kp, str)][:4],
                        retrieval_qas=[qa for qa in obj.get('retrieval_qas', []) if isinstance(qa, dict) and 'q' in qa and 'a' in qa][:2],
                    )
                )
        except Exception:
            continue
    # Ensure every segment gets a summary.
    have = {s.segment_index for s in summaries}
    for s in segments:
        if s.index not in have:
            summaries.append(
                SegmentSummary(
                    segment_index=s.index,
                    micro_summary=_fallback_micro_summary(s.text),
                    key_points=[],
                    retrieval_qas=[],
                )
            )
    # Sort by segment index
    summaries.sort(key=lambda x: x.segment_index)
    return summaries
