from __future__ import annotations

import json
import re
from typing import Any, Dict, List, Optional, Sequence, Tuple

from ..common.utils import ChatGPT_API, count_tokens, extract_json
from ..prompts.transcript import render_transcript_topics_prompt


TIMESTAMP_RE = re.compile(r"\b(\d{1,2}:\d{2}(?::\d{2})?)\b")


def _time_to_seconds(value: Optional[str]) -> Optional[int]:
    if not value:
        return None
    raw = str(value).strip()
    if not raw:
        return None
    m = TIMESTAMP_RE.search(raw)
    if not m:
        return None
    parts = [int(x) for x in m.group(1).split(":")]
    if len(parts) == 2:
        mm, ss = parts
        return mm * 60 + ss
    if len(parts) == 3:
        hh, mm, ss = parts
        return hh * 3600 + mm * 60 + ss
    return None


def _seconds_to_time(value: Optional[int]) -> Optional[str]:
    if value is None:
        return None
    if value < 0:
        return None
    hh = value // 3600
    rem = value % 3600
    mm = rem // 60
    ss = rem % 60
    if hh > 0:
        return f"{hh:02d}:{mm:02d}:{ss:02d}"
    return f"{mm:02d}:{ss:02d}"


def _extract_time_bounds(text: str) -> Tuple[Optional[str], Optional[str]]:
    matches = [m.group(1) for m in TIMESTAMP_RE.finditer(text or "")]
    if not matches:
        return None, None
    start = matches[0]
    end = matches[-1]
    return start, end


def _build_raw_text(page_text_by_page: Dict[int, str]) -> Tuple[str, List[Dict[str, int]]]:
    pages = sorted(page_text_by_page)
    parts: List[str] = []
    page_spans: List[Dict[str, int]] = []
    offset = 0

    for page in pages:
        text = (page_text_by_page.get(page) or "").strip()
        if not text:
            continue
        if parts:
            parts.append("\n\n")
            offset += 2
        start = offset
        parts.append(text)
        offset += len(text)
        end_exclusive = offset
        page_spans.append(
            {
                "page": int(page),
                "start_char": int(start),
                "end_char_exclusive": int(end_exclusive),
            }
        )

    raw_text = "".join(parts)
    return raw_text, page_spans


def _char_to_page(page_spans: Sequence[Dict[str, int]], char_index: int) -> Optional[int]:
    for row in page_spans:
        if row["start_char"] <= char_index < row["end_char_exclusive"]:
            return row["page"]
    if page_spans and char_index >= page_spans[-1]["end_char_exclusive"]:
        return int(page_spans[-1]["page"])
    return int(page_spans[0]["page"]) if page_spans else None


def _paragraph_spans(raw_text: str) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []
    for m in re.finditer(r"\S[\s\S]*?(?=\n\s*\n|$)", raw_text):
        start, end = m.span()
        text = m.group(0).strip()
        if not text:
            continue
        out.append({"start_char": start, "end_char_exclusive": end, "text": text})
    return out


def _build_chunks(
    raw_text: str,
    page_spans: Sequence[Dict[str, int]],
    *,
    max_token_num_each_node: int,
    max_page_num_each_node: int,
) -> List[Dict[str, Any]]:
    paragraphs = _paragraph_spans(raw_text)
    if not paragraphs and raw_text.strip():
        paragraphs = [
            {
                "start_char": 0,
                "end_char_exclusive": len(raw_text),
                "text": raw_text.strip(),
            }
        ]

    target_tokens = max(300, min(1500, max(1, max_token_num_each_node) // 4))
    out: List[Dict[str, Any]] = []
    buf: List[Dict[str, Any]] = []
    buf_tokens = 0

    def flush() -> None:
        nonlocal buf, buf_tokens
        if not buf:
            return
        chunk_text = "\n\n".join(p["text"] for p in buf).strip()
        start_char = int(buf[0]["start_char"])
        end_exclusive = int(buf[-1]["end_char_exclusive"])
        end_char = max(start_char, end_exclusive - 1)
        start_page = _char_to_page(page_spans, start_char) or 1
        end_page = _char_to_page(page_spans, end_char) or start_page
        start_time, end_time = _extract_time_bounds(chunk_text)
        out.append(
            {
                "node_id": f"N{len(out) + 1:04d}",
                "node_index": len(out) + 1,
                "start_char": start_char,
                "end_char": end_char,
                "text": chunk_text,
                "start_page": int(start_page),
                "end_page": int(end_page),
                "start_time": start_time,
                "end_time": end_time,
            }
        )
        buf = []
        buf_tokens = 0

    for para in paragraphs:
        para_tokens = count_tokens(para["text"], model="gpt-4o")
        para_start_page = _char_to_page(page_spans, int(para["start_char"])) or 1
        para_end_char = max(int(para["start_char"]), int(para["end_char_exclusive"]) - 1)
        para_end_page = _char_to_page(page_spans, para_end_char) or para_start_page

        exceeds_token = buf and (buf_tokens + para_tokens > target_tokens)
        if buf:
            buf_start_page = _char_to_page(page_spans, int(buf[0]["start_char"])) or para_start_page
            exceeds_page = (para_end_page - buf_start_page + 1) > max(1, max_page_num_each_node)
        else:
            exceeds_page = False

        if exceeds_token or exceeds_page:
            flush()

        buf.append(para)
        buf_tokens += para_tokens

    flush()
    return out


def _render_chunks_for_prompt(chunks: Sequence[Dict[str, Any]]) -> str:
    slim = []
    for c in chunks:
        slim.append(
            {
                "node_id": c["node_id"],
                "start_char": c["start_char"],
                "end_char": c["end_char"],
                "start_time": c.get("start_time"),
                "end_time": c.get("end_time"),
                "text": c["text"],
            }
        )
    return json.dumps(slim, ensure_ascii=False, indent=2)


def _parse_topic_groups(raw: str) -> List[Dict[str, Any]]:
    parsed = extract_json(raw)
    if isinstance(parsed, dict) and isinstance(parsed.get("topics"), list):
        parsed = parsed["topics"]
    if not isinstance(parsed, list):
        try:
            reparsed = json.loads(raw)
            if isinstance(reparsed, list):
                parsed = reparsed
            elif isinstance(reparsed, dict) and isinstance(reparsed.get("topics"), list):
                parsed = reparsed["topics"]
        except Exception:
            parsed = []

    out: List[Dict[str, Any]] = []
    if not isinstance(parsed, list):
        return out

    for item in parsed:
        if not isinstance(item, dict):
            continue
        topic = re.sub(r"\s+", " ", str(item.get("topic") or item.get("heading") or "").strip())
        if not topic:
            continue
        ids_raw = item.get("node_ids")
        if not isinstance(ids_raw, list):
            # Backward compatibility with earlier prompt versions.
            ids_raw = item.get("chunk_ids")
        if not isinstance(ids_raw, list):
            ids_raw = []
        node_ids = [str(x).strip() for x in ids_raw if str(x).strip()]
        out.append(
            {
                "topic": topic,
                "node_ids": node_ids,
                "start_time": str(item.get("start_time", "")).strip() or None,
                "end_time": str(item.get("end_time", "")).strip() or None,
                "summary": re.sub(r"\s+", " ", str(item.get("summary", "")).strip()),
                "sentiment": str(item.get("sentiment", "")).strip() or None,
            }
        )
    return out


def _finalize_topic_groups(
    groups: Sequence[Dict[str, Any]],
    chunks: Sequence[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    by_id = {str(c["node_id"]): c for c in chunks}
    ordered_ids = [str(c["node_id"]) for c in chunks]
    assigned: set[str] = set()
    out: List[Dict[str, Any]] = []

    for g in groups:
        unique_ids: List[str] = []
        for nid in g.get("node_ids", []):
            if nid in by_id and nid not in assigned:
                unique_ids.append(nid)
                assigned.add(nid)
        if not unique_ids:
            continue
        out.append({**g, "node_ids": unique_ids})

    unassigned = [cid for cid in ordered_ids if cid not in assigned]
    if unassigned:
        out.append(
            {
                "topic": "Other Transcript Discussion",
                "node_ids": unassigned,
                "start_time": None,
                "end_time": None,
                "summary": "",
                "sentiment": None,
            }
        )
    return out


def build_transcript_sections(
    *,
    page_text_by_page: Dict[int, str],
    model: str,
    max_token_num_each_node: int,
    max_page_num_each_node: int,
    prompt_template: Optional[str] = None,
    include_topic_summary: bool = True,
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    raw_text, page_spans = _build_raw_text(page_text_by_page)
    chunks = _build_chunks(
        raw_text,
        page_spans,
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
    )

    topics_raw: List[Dict[str, Any]] = []
    if chunks:
        chunk_payload = _render_chunks_for_prompt(chunks)
        prompt = render_transcript_topics_prompt(content=chunk_payload, template=prompt_template)
        raw = ChatGPT_API(model=model, prompt=prompt)
        topics_raw = [] if raw == "Error" else _parse_topic_groups(str(raw))

    grouped_topics = _finalize_topic_groups(topics_raw, chunks)
    chunk_map = {str(c["node_id"]): c for c in chunks}

    sections: List[Dict[str, Any]] = []
    transcript_topics_payload: List[Dict[str, Any]] = []

    for topic_idx, topic in enumerate(grouped_topics, start=1):
        node_ids = topic["node_ids"]
        topic_chunks = [chunk_map[nid] for nid in node_ids if nid in chunk_map]
        if not topic_chunks:
            continue

        topic_start_char = min(int(c["start_char"]) for c in topic_chunks)
        topic_end_char = max(int(c["end_char"]) for c in topic_chunks)
        topic_start_page = min(int(c["start_page"]) for c in topic_chunks)
        topic_end_page = max(int(c["end_page"]) for c in topic_chunks)

        first_chunk = topic_chunks[0]
        last_chunk = topic_chunks[-1]
        start_time = topic.get("start_time") or first_chunk.get("start_time")
        end_time = topic.get("end_time") or last_chunk.get("end_time")
        if not start_time or not end_time:
            s_sec = _time_to_seconds(start_time) if start_time else None
            e_sec = _time_to_seconds(end_time) if end_time else None
            if s_sec is None:
                s_sec = _time_to_seconds(first_chunk.get("start_time"))
            if e_sec is None:
                e_sec = _time_to_seconds(last_chunk.get("end_time"))
            start_time = _seconds_to_time(s_sec)
            end_time = _seconds_to_time(e_sec)

        topic_id = f"T{topic_idx:02d}"
        child_sections: List[Dict[str, Any]] = []
        topic_text_parts: List[str] = []
        topic_chunks_payload: List[Dict[str, Any]] = []

        for child_idx, chunk in enumerate(topic_chunks, start=1):
            chunk_text = str(chunk.get("text", "")).strip()
            topic_text_parts.append(chunk_text)
            chunk_id = f"{topic_id}_C{child_idx:02d}"
            child_sections.append(
                {
                    "title": f"{topic.get('topic', 'Topic')} - Node {child_idx}",
                    "text": chunk_text,
                    "start_index": int(chunk["start_page"]),
                    "end_index": int(chunk["end_page"]),
                    "level": 2,
                    "start_char": int(chunk["start_char"]),
                    "end_char": int(chunk["end_char"]),
                    "topic_meta": {
                        "topic_id": topic_id,
                        "node_id": chunk_id,
                        "source_node_id": chunk["node_id"],
                        "start_time": chunk.get("start_time"),
                        "end_time": chunk.get("end_time"),
                    },
                }
            )
            chunk_payload: Dict[str, Any] = {
                "node_id": chunk_id,
                "span": {
                    "start_char": int(chunk["start_char"]),
                    "end_char": int(chunk["end_char"]),
                },
                "start_time": chunk.get("start_time"),
                "end_time": chunk.get("end_time"),
                "text": chunk_text,
            }
            if include_topic_summary:
                chunk_payload["summary"] = ""
            topic_chunks_payload.append(chunk_payload)

        section_meta: Dict[str, Any] = {
            "topic_id": topic_id,
            "start_time": start_time,
            "end_time": end_time,
            "start_char": topic_start_char,
            "end_char": topic_end_char,
            "sentiment": topic.get("sentiment"),
        }
        if include_topic_summary:
            section_meta["summary"] = topic.get("summary") or ""

        sections.append(
            {
                "title": topic.get("topic", f"Topic {topic_idx}"),
                "text": "",
                "start_index": topic_start_page,
                "end_index": topic_end_page,
                "level": 1,
                "start_char": topic_start_char,
                "end_char": topic_end_char,
                "topic_meta": section_meta,
                "children": child_sections,
            }
        )

        payload_topic: Dict[str, Any] = {
            "topic_id": topic_id,
            "title": topic.get("topic", f"Topic {topic_idx}"),
            "span": {"start_char": topic_start_char, "end_char": topic_end_char},
            "start_time": start_time,
            "end_time": end_time,
            "full_text": "\n\n".join(part for part in topic_text_parts if part),
            "nodes": topic_chunks_payload,
            "sentiment": {"label": topic.get("sentiment") or "Neutral"},
        }
        if include_topic_summary:
            payload_topic["summary"] = topic.get("summary") or ""
        transcript_topics_payload.append(payload_topic)

    meta = {
        "mode": "transcript_topic_sections",
        "table_of_contents_pages": [],
        "table_of_contents_entries": [
            {
                "structure": str(i),
                "title": row["title"],
                "page": None,
                "physical_index": row["start_index"],
            }
            for i, row in enumerate(sections, start=1)
        ],
        "verify_accuracy": 1.0 if sections else 0.0,
        "verify_incorrect_count": 0,
        "llm_pipeline_enabled": True,
        "steps": ["segment_transcript_nodes", "group_nodes_by_topic"],
        "generated_topic_count": len(sections),
        "raw_text_length": len(raw_text),
        "transcript": {
            "raw_text_length": len(raw_text),
            "topics": transcript_topics_payload,
        },
    }
    return sections, meta
