from __future__ import annotations

import json
import math
import re
from typing import Any, Dict, List, Optional

from ..common.utils import ChatGPT_API, extract_json
from ..prompts.semi_structured import render_semi_structured_headings_prompt


def _norm(text: str) -> str:
    return re.sub(r"\s+", " ", (text or "")).strip()


def _norm_key(text: str) -> str:
    return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9]+", " ", (text or "").lower())).strip()


def _dedupe_entries(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []
    seen: set[str] = set()
    for item in entries:
        title = _norm(str(item.get("title", "")))
        if not title:
            continue
        key = _norm_key(title)
        if not key or key in seen:
            continue
        seen.add(key)
        out.append(
            {
                "structure": item.get("structure"),
                "title": title,
                "page": None,
                "physical_index": item.get("physical_index") if isinstance(item.get("physical_index"), int) else None,
            }
        )
    return out


def _page_groups(page_text_by_page: Dict[int, str], max_pages_per_group: int = 8) -> List[str]:
    pages = sorted(page_text_by_page)
    if not pages:
        return []
    groups: List[str] = []
    for i in range(0, len(pages), max_pages_per_group):
        page_slice = pages[i : i + max_pages_per_group]
        wrapped = []
        for page in page_slice:
            wrapped.append(f"<physical_index_{page}>\n{page_text_by_page.get(page, '')}\n<physical_index_{page}>")
        groups.append("\n\n".join(wrapped))
    return groups


def _extract_physical_index(value: Any) -> Optional[int]:
    if isinstance(value, int):
        return value
    if isinstance(value, str):
        match = re.search(r"physical_index_(\d+)", value)
        if match:
            return int(match.group(1))
        if value.strip().isdigit():
            return int(value.strip())
    return None


def _heuristic_title(page_text: str, fallback: str) -> str:
    lines = [ln.strip() for ln in (page_text or "").splitlines() if ln.strip()]
    for line in lines[:20]:
        if len(line) < 5 or len(line) > 120:
            continue
        if re.match(r"^\d{1,2}:\d{2}$", line):
            continue
        if re.search(r"\b(reference id|started transcription)\b", line, flags=re.IGNORECASE):
            continue
        if re.match(r"^\d+\s+of\s+\d+$", line, flags=re.IGNORECASE):
            continue
        words = line.split()
        if 2 <= len(words) <= 14:
            return line
    return fallback


def _heuristic_entries(page_text_by_page: Dict[int, str]) -> List[Dict[str, Any]]:
    pages = sorted(page_text_by_page)
    if not pages:
        return []
    target_count = max(2, min(10, math.ceil(len(pages) / 3)))
    step = max(1, math.ceil(len(pages) / target_count))
    entries: List[Dict[str, Any]] = []
    section_num = 1
    for idx in range(0, len(pages), step):
        page = pages[idx]
        fallback = f"Section {section_num}"
        title = _heuristic_title(page_text_by_page.get(page, ""), fallback=fallback)
        entries.append(
            {
                "structure": str(section_num),
                "title": title,
                "page": None,
                "physical_index": page,
            }
        )
        section_num += 1
    return _dedupe_entries(entries)


def _parse_llm_entries(raw: str) -> List[Dict[str, Any]]:
    parsed = extract_json(raw)
    if isinstance(parsed, dict) and isinstance(parsed.get("entries"), list):
        parsed = parsed["entries"]
    if not isinstance(parsed, list):
        try:
            reparsed = json.loads(raw)
            if isinstance(reparsed, list):
                parsed = reparsed
            elif isinstance(reparsed, dict) and isinstance(reparsed.get("entries"), list):
                parsed = reparsed["entries"]
        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
        title = _norm(str(item.get("title", "")))
        if not title:
            continue
        out.append(
            {
                "structure": _norm(str(item.get("structure", ""))) or None,
                "title": title,
                "page": None,
                "physical_index": _extract_physical_index(item.get("physical_index")),
            }
        )
    return _dedupe_entries(out)


def _generate_entries_with_llm(
    page_text_by_page: Dict[int, str],
    model: Optional[str],
    prompt_template: Optional[str] = None,
) -> List[Dict[str, Any]]:
    if not model:
        return []

    groups = _page_groups(page_text_by_page, max_pages_per_group=8)
    generated: List[Dict[str, Any]] = []
    for group in groups:
        prompt = render_semi_structured_headings_prompt(content=group, template=prompt_template)
        raw = ChatGPT_API(model=model, prompt=prompt)
        if raw == "Error":
            continue
        generated.extend(_parse_llm_entries(raw))

    generated = _dedupe_entries(generated)
    generated.sort(key=lambda x: x.get("physical_index") if isinstance(x.get("physical_index"), int) else 10**9)
    return generated


def _normalize_existing_entries(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    normalized = _dedupe_entries(entries)
    normalized.sort(key=lambda x: x.get("physical_index") if isinstance(x.get("physical_index"), int) else 10**9)
    return normalized


def build_semi_structured_outline(
    *,
    base_outline: Dict[str, Any],
    page_text_by_page: Dict[int, str],
    model: Optional[str],
    use_llm: bool,
    prompt_template: Optional[str] = None,
) -> Dict[str, Any]:
    entries = _normalize_existing_entries(base_outline.get("table_of_contents_entries", []) or [])
    weak_outline = len(entries) < 2 or all(item.get("physical_index") is None for item in entries)

    generated: List[Dict[str, Any]] = []
    source = "existing"
    if weak_outline and use_llm:
        generated = _generate_entries_with_llm(page_text_by_page=page_text_by_page, model=model, prompt_template=prompt_template)
        if generated:
            source = "llm"
    if weak_outline and not generated:
        generated = _heuristic_entries(page_text_by_page)
        if generated:
            source = "heuristic"

    final_entries = generated if generated else entries
    final = dict(base_outline)
    final["mode"] = "semi_structured_enhanced"
    final["table_of_contents_entries"] = final_entries
    final["steps"] = [*(base_outline.get("steps") or []), "semi_structured_heading_normalization"]
    final["semi_structured_source"] = source
    if generated and source != "existing":
        final["verify_accuracy"] = 1.0 if source == "llm" else float(base_outline.get("verify_accuracy", 0.0))
        final["verify_incorrect_count"] = 0 if source == "llm" else int(base_outline.get("verify_incorrect_count", 0))
    return final
