from __future__ import annotations

import logging
import math
import os
import re
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

from ..common.utils import ChatGPT_API, count_tokens, extract_json
from ..llm.client import has_llm_credentials

MODE_TABLE_OF_CONTENTS_WITH_PAGE_NUMBERS = "table_of_contents_with_page_numbers"
MODE_TABLE_OF_CONTENTS_WITHOUT_PAGE_NUMBERS = "table_of_contents_without_page_numbers"
MODE_SYNTHETIC_TABLE_OF_CONTENTS = "synthetic_table_of_contents"
MODE_FALLBACK_SYNTHETIC_TABLE_OF_CONTENTS = "fallback_synthetic_table_of_contents"


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


def _llm_enabled(model: Optional[str], use_llm: bool = True) -> bool:
    if not use_llm:
        return False
    if os.getenv("NAVEXA_DISABLE_LLM_PIPELINE", "0").strip().lower() in {"1", "true", "yes", "on"}:
        return False
    if not model:
        return False
    return has_llm_credentials()


def _unique_by_title(items: Iterable[Dict[str, Any]]) -> List[Dict[str, Any]]:
    seen: set[str] = set()
    out: List[Dict[str, Any]] = []
    for item in items:
        key = _norm(str(item.get("title", "")))
        if not key or key in seen:
            continue
        seen.add(key)
        out.append(item)
    return out


def _is_numeric_heading(line: str) -> bool:
    return bool(re.match(r"^\s*\d+(?:\.\d+)*\s+\S", line or ""))


def _looks_like_toc_page(page_text: str) -> bool:
    if not page_text:
        return False
    lines = [ln.strip() for ln in page_text.splitlines() if ln.strip()]
    if not lines:
        return False

    contains_contents = bool(re.search(r"\b(table of contents|contents)\b", page_text, re.IGNORECASE))
    numeric_lines = sum(1 for ln in lines if _is_numeric_heading(ln))
    dotted_lines = sum(1 for ln in lines if re.search(r"\.{2,}\s*\d+\s*$", ln))
    short_lines = sum(1 for ln in lines if len(ln) < 120)

    if contains_contents and (numeric_lines >= 3 or dotted_lines >= 2):
        return True
    if numeric_lines >= 10 and short_lines >= max(5, int(len(lines) * 0.4)):
        return True
    return False


def toc_detector_single_page(content: str, model: Optional[str] = None, use_llm: bool = True) -> str:
    if _llm_enabled(model, use_llm=use_llm):
        prompt = f"""
You are given one page of PDF text. Detect whether this page is a table of contents page.

Rules:
- Return "yes" only when this page is a TOC/index page listing sections.
- Abstract/list of figures/list of tables are not TOC.

Reply as JSON:
{{
  "toc_detected": "yes or no"
}}

Page text:
{content}
"""
        try:
            raw = ChatGPT_API(model=model, prompt=prompt)
            parsed = extract_json(raw)
            detected = str(parsed.get("toc_detected", "")).strip().lower()
            if detected in {"yes", "no"}:
                return detected
        except Exception:
            pass

    return "yes" if _looks_like_toc_page(content) else "no"


def find_toc_pages(
    page_text_by_page: Dict[int, str],
    toc_check_page_num: int = 20,
    model: Optional[str] = None,
    use_llm: bool = True,
) -> List[int]:
    toc_pages: List[int] = []
    last_was_toc = False

    for page_no in sorted(page_text_by_page):
        if page_no > toc_check_page_num and not last_was_toc:
            break
        detected = toc_detector_single_page(page_text_by_page[page_no], model=model, use_llm=use_llm)
        if detected == "yes":
            toc_pages.append(page_no)
            last_was_toc = True
        elif last_was_toc:
            break

    return toc_pages


def _parse_toc_line(line: str) -> Optional[Dict[str, Any]]:
    ln = re.sub(r"\s+", " ", (line or "").strip())
    if not ln or len(ln) > 220:
        return None
    if re.match(r"^\d+\s+of\s+\d+$", ln, flags=re.IGNORECASE):
        return None
    if ln.lower().startswith("reference id"):
        return None

    numbered = re.match(r"^((?:\d+\.)*\d+)\s+(.+?)(?:\s+(\d{1,4}))?$", ln)
    if numbered:
        structure = numbered.group(1)
        suffix = numbered.group(2).strip()
        page = numbered.group(3)
        return {
            "structure": structure,
            "title": f"{structure} {suffix}".strip(),
            "page": int(page) if page and page.isdigit() else None,
            "physical_index": None,
        }

    return None


def _heuristic_extract_toc_entries(toc_text: str) -> List[Dict[str, Any]]:
    entries = [_parse_toc_line(line) for line in toc_text.splitlines()]
    return _unique_by_title([e for e in entries if e is not None])


def toc_transformer(toc_content: str, model: Optional[str] = None, use_llm: bool = True) -> List[Dict[str, Any]]:
    if _llm_enabled(model, use_llm=use_llm):
        prompt = f"""
Convert raw TOC text into JSON.

Return ONLY a JSON list with objects:
[
  {{
    "structure": "1 or 1.2 etc",
    "title": "full section title",
    "page": <integer or null>
  }}
]

Raw TOC text:
{toc_content}
"""
        try:
            raw = ChatGPT_API(model=model, prompt=prompt)
            parsed = extract_json(raw)
            if isinstance(parsed, list):
                out: List[Dict[str, Any]] = []
                for item in parsed:
                    if not isinstance(item, dict):
                        continue
                    title = str(item.get("title", "")).strip()
                    structure = item.get("structure")
                    page = item.get("page")
                    if not title:
                        continue
                    page_int = int(page) if isinstance(page, int) or (isinstance(page, str) and page.isdigit()) else None
                    out.append(
                        {
                            "structure": str(structure).strip() if structure is not None else None,
                            "title": title,
                            "page": page_int,
                            "physical_index": None,
                        }
                    )
                out = _unique_by_title(out)
                if out:
                    return out
        except Exception:
            pass

    return _heuristic_extract_toc_entries(toc_content)


def _title_token_overlap(title: str, page_text: str) -> int:
    a = set(_norm(title).split())
    b = set(_norm(page_text).split())
    if not a or not b:
        return 0
    return len(a & b)


def _match_title_to_page(
    title: str,
    page_text_by_page: Dict[int, str],
    start_page: int = 1,
    disallow_pages: Optional[set[int]] = None,
) -> Optional[int]:
    disallow = disallow_pages or set()
    target = _norm(title)
    if not target:
        return None

    best_page: Optional[int] = None
    best_score = -10_000
    for page_no in sorted(page_text_by_page):
        if page_no < start_page or page_no in disallow:
            continue
        text = page_text_by_page[page_no]
        text_norm = _norm(text)
        if not text_norm:
            continue

        score = 0
        if target in text_norm:
            score += 10
        overlap = _title_token_overlap(title, text)
        score += overlap
        if overlap < 2:
            score -= 5

        if score > best_score:
            best_score = score
            best_page = page_no
        if score >= 12:
            return page_no

    if best_score >= 5:
        return best_page
    return None


def add_page_number_to_toc(
    entries: List[Dict[str, Any]],
    page_text_by_page: Dict[int, str],
    toc_pages: Sequence[int],
) -> List[Dict[str, Any]]:
    disallow = set(toc_pages)
    cursor = max(disallow) + 1 if disallow else 1
    out: List[Dict[str, Any]] = []
    for item in entries:
        candidate = dict(item)
        matched = _match_title_to_page(candidate.get("title", ""), page_text_by_page, start_page=cursor, disallow_pages=disallow)
        if matched is None:
            matched = _match_title_to_page(candidate.get("title", ""), page_text_by_page, start_page=1, disallow_pages=disallow)
        candidate["physical_index"] = matched
        if matched is not None:
            cursor = max(cursor, matched)
        out.append(candidate)
    return out


def process_toc_no_page_numbers(
    toc_entries: List[Dict[str, Any]],
    page_text_by_page: Dict[int, str],
    toc_pages: Sequence[int],
    model: Optional[str] = None,
) -> List[Dict[str, Any]]:
    del model  # parity placeholder: PageIndex has LLM-heavy flow; mapping here is deterministic + verified later.
    return add_page_number_to_toc(toc_entries, page_text_by_page, toc_pages)


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


def _page_groups(page_text_by_page: Dict[int, str], max_tokens: int = 18_000, overlap_pages: int = 1) -> List[str]:
    pages = sorted(page_text_by_page)
    chunks: List[str] = []
    buf: List[Tuple[int, str, int]] = []
    current_tokens = 0

    for p in pages:
        text = page_text_by_page[p]
        wrapped = f"<physical_index_{p}>\n{text}\n<physical_index_{p}>\n"
        tok = count_tokens(wrapped, model="gpt-4o")
        if current_tokens + tok > max_tokens and buf:
            chunks.append("\n".join(x[1] for x in buf))
            overlap = buf[-overlap_pages:] if overlap_pages > 0 else []
            buf = overlap.copy()
            current_tokens = sum(x[2] for x in buf)
        buf.append((p, wrapped, tok))
        current_tokens += tok

    if buf:
        chunks.append("\n".join(x[1] for x in buf))
    return chunks


def generate_toc_init(part: str, model: Optional[str] = None, use_llm: bool = True) -> List[Dict[str, Any]]:
    if not _llm_enabled(model, use_llm=use_llm):
        return []
    prompt = f"""
Extract document section hierarchy from this partial PDF text.

Return ONLY a JSON list:
[
  {{
    "structure": "1 / 1.1 / 2.3.4",
    "title": "section title exactly from document",
    "physical_index": "<physical_index_X>"
  }}
]

Document text:
{part}
"""
    raw = ChatGPT_API(model=model, prompt=prompt)
    parsed = extract_json(raw)
    if not isinstance(parsed, list):
        return []
    out: List[Dict[str, Any]] = []
    for item in parsed:
        if not isinstance(item, dict):
            continue
        title = str(item.get("title", "")).strip()
        if not title:
            continue
        out.append(
            {
                "structure": str(item.get("structure", "")).strip() or None,
                "title": title,
                "page": None,
                "physical_index": _extract_physical_from_tag(item.get("physical_index")),
            }
        )
    return _unique_by_title(out)


def generate_toc_continue(
    toc_content: List[Dict[str, Any]],
    part: str,
    model: Optional[str] = None,
    use_llm: bool = True,
) -> List[Dict[str, Any]]:
    if not _llm_enabled(model, use_llm=use_llm):
        return []
    prompt = f"""
Continue TOC extraction for next document chunk.

Return ONLY additional items as JSON list:
[
  {{
    "structure": "1 / 1.1 / 2.3.4",
    "title": "section title",
    "physical_index": "<physical_index_X>"
  }}
]

Existing TOC:
{toc_content}

New chunk:
{part}
"""
    raw = ChatGPT_API(model=model, prompt=prompt)
    parsed = extract_json(raw)
    if not isinstance(parsed, list):
        return []
    out: List[Dict[str, Any]] = []
    for item in parsed:
        if not isinstance(item, dict):
            continue
        title = str(item.get("title", "")).strip()
        if not title:
            continue
        out.append(
            {
                "structure": str(item.get("structure", "")).strip() or None,
                "title": title,
                "page": None,
                "physical_index": _extract_physical_from_tag(item.get("physical_index")),
            }
        )
    return _unique_by_title(out)


def _fallback_synthetic_toc(markdown_nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    items: List[Dict[str, Any]] = []
    for node in markdown_nodes:
        title = str(node.get("title", "")).strip()
        if not title:
            continue
        m = re.match(r"^\s*(\d+(?:\.\d+)*)\s+(.+)$", title)
        if not m:
            continue
        items.append({"structure": m.group(1), "title": title, "page": None, "physical_index": None})
    return _unique_by_title(items)


def process_no_toc(
    page_text_by_page: Dict[int, str],
    markdown_nodes: List[Dict[str, Any]],
    model: Optional[str] = None,
    use_llm: bool = True,
) -> List[Dict[str, Any]]:
    groups = _page_groups(page_text_by_page)
    generated: List[Dict[str, Any]] = []

    if groups and _llm_enabled(model, use_llm=use_llm):
        generated = generate_toc_init(groups[0], model=model, use_llm=use_llm)
        for part in groups[1:]:
            generated.extend(generate_toc_continue(generated, part, model=model, use_llm=use_llm))
        generated = _unique_by_title(generated)

    if not generated:
        generated = _fallback_synthetic_toc(markdown_nodes)

    if not generated:
        return []

    missing_indexes = any(item.get("physical_index") is None for item in generated)
    if missing_indexes:
        generated = add_page_number_to_toc(generated, page_text_by_page, toc_pages=[])

    return generated


def process_toc_with_page_numbers(
    toc_entries: List[Dict[str, Any]],
    page_text_by_page: Dict[int, str],
    toc_pages: Sequence[int],
    model: Optional[str] = None,
) -> List[Dict[str, Any]]:
    del model  # parity placeholder
    if not toc_entries:
        return []

    # Estimate page-number offset by matching first few headings.
    offsets: List[int] = []
    for item in toc_entries[:20]:
        page_number = item.get("page")
        title = item.get("title", "")
        if not isinstance(page_number, int):
            continue
        matched = _match_title_to_page(title, page_text_by_page, start_page=1, disallow_pages=set(toc_pages))
        if matched is not None:
            offsets.append(matched - page_number)

    offset = 0
    if offsets:
        counts: Dict[int, int] = {}
        for off in offsets:
            counts[off] = counts.get(off, 0) + 1
        offset = max(counts.items(), key=lambda x: x[1])[0]

    out: List[Dict[str, Any]] = []
    for item in toc_entries:
        candidate = dict(item)
        page_number = candidate.get("page")
        if isinstance(page_number, int):
            candidate["physical_index"] = page_number + offset
        else:
            candidate["physical_index"] = None
        out.append(candidate)

    return add_page_number_to_toc(out, page_text_by_page, toc_pages=toc_pages)


def _title_on_page(title: str, page_text: str) -> bool:
    t = _norm(title)
    p = _norm(page_text)
    if not t or not p:
        return False
    if t in p:
        return True
    return _title_token_overlap(title, page_text) >= max(2, math.ceil(len(t.split()) * 0.45))


def verify_toc(
    page_text_by_page: Dict[int, str],
    list_result: List[Dict[str, Any]],
    model: Optional[str] = None,
) -> Tuple[float, List[Dict[str, Any]]]:
    del model  # parity placeholder: deterministic verifier keeps pipeline stable without API.
    checked = 0
    correct = 0
    incorrect: List[Dict[str, Any]] = []

    for idx, item in enumerate(list_result):
        page_no = item.get("physical_index")
        title = str(item.get("title", "")).strip()
        if not isinstance(page_no, int) or page_no not in page_text_by_page or not title:
            continue
        checked += 1
        if _title_on_page(title, page_text_by_page[page_no]):
            correct += 1
        else:
            incorrect.append({"list_index": idx, "title": title, "physical_index": page_no})

    if checked == 0:
        return 0.0, []
    return correct / checked, incorrect


def _fix_incorrect_toc(
    toc_entries: List[Dict[str, Any]],
    incorrect_results: List[Dict[str, Any]],
    page_text_by_page: Dict[int, str],
    toc_pages: Sequence[int],
) -> List[Dict[str, Any]]:
    fixed = [dict(item) for item in toc_entries]
    disallow = set(toc_pages)
    for miss in incorrect_results:
        idx = miss.get("list_index")
        if not isinstance(idx, int) or idx < 0 or idx >= len(fixed):
            continue
        title = fixed[idx].get("title", "")
        prev_page = 1
        next_page = max(page_text_by_page) if page_text_by_page else 1
        for j in range(idx - 1, -1, -1):
            p = fixed[j].get("physical_index")
            if isinstance(p, int):
                prev_page = p
                break
        for j in range(idx + 1, len(fixed)):
            p = fixed[j].get("physical_index")
            if isinstance(p, int):
                next_page = p
                break
        local_pages = {k: v for k, v in page_text_by_page.items() if prev_page <= k <= next_page}
        matched = _match_title_to_page(title, local_pages, start_page=prev_page, disallow_pages=disallow)
        if matched is not None:
            fixed[idx]["physical_index"] = matched
    return fixed


def _flatten_markdown_tree(markdown_tree: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []

    def walk(nodes: List[Dict[str, Any]]) -> None:
        for node in nodes:
            out.append({"title": node.get("title", ""), "level": node.get("level", 1), "text": node.get("text", "")})
            walk(node.get("nodes") or [])

    walk(markdown_tree)
    return out


@dataclass
class OutlineResult:
    mode: str
    toc_pages: List[int]
    toc_entries: List[Dict[str, Any]]
    verify_accuracy: float
    verify_incorrect_count: int
    llm_pipeline_enabled: bool
    steps: List[str]

    def as_dict(self) -> Dict[str, Any]:
        return {
            "mode": self.mode,
            "table_of_contents_pages": self.toc_pages,
            "table_of_contents_entries": self.toc_entries,
            "verify_accuracy": round(self.verify_accuracy, 4),
            "verify_incorrect_count": self.verify_incorrect_count,
            "llm_pipeline_enabled": self.llm_pipeline_enabled,
            "steps": self.steps,
        }


def build_outline_engine(
    markdown_tree: List[Dict[str, Any]],
    page_text_by_page: Dict[int, str],
    model: Optional[str] = None,
    toc_check_page_num: int = 20,
    use_llm: bool = True,
    logger: Optional[logging.Logger] = None,
) -> Dict[str, Any]:
    steps: List[str] = []
    flat_md = _flatten_markdown_tree(markdown_tree)
    llm_enabled = _llm_enabled(model, use_llm=use_llm)

    if logger and logger.isEnabledFor(logging.DEBUG):
        logger.debug(
            "[table-of-contents] strategy=%s pages_considered=%d markdown_headings=%d",
            "llm-assisted" if llm_enabled else "heuristic",
            len(page_text_by_page),
            len(flat_md),
        )

    toc_pages = find_toc_pages(
        page_text_by_page,
        toc_check_page_num=toc_check_page_num,
        model=model,
        use_llm=use_llm,
    )
    steps.append("detect_table_of_contents_pages")
    if logger:
        logger.info("[table-of-contents] detected pages: %s", toc_pages if toc_pages else "none")

    if toc_pages:
        toc_text = "\n".join(page_text_by_page[p] for p in toc_pages if p in page_text_by_page)
        toc_entries = toc_transformer(toc_text, model=model, use_llm=use_llm)
        steps.append("transform_table_of_contents_text")
        if logger:
            logger.info("[table-of-contents] parsed entries: %d", len(toc_entries))
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(
                    "[table-of-contents] preview=%s",
                    [e.get("title", "") for e in toc_entries[:6]],
                )
                logger.debug("[table-of-contents] source_chars=%d", len(toc_text))
        has_page_numbers = any(isinstance(item.get("page"), int) for item in toc_entries)
        if logger and logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                "[table-of-contents] mapping_strategy=%s",
                "existing_page_numbers" if has_page_numbers else "title_to_page_matching",
            )
        if has_page_numbers:
            toc_entries = process_toc_with_page_numbers(toc_entries, page_text_by_page, toc_pages=toc_pages, model=model)
            mode = MODE_TABLE_OF_CONTENTS_WITH_PAGE_NUMBERS
            steps.append("map_table_of_contents_using_existing_page_numbers")
        else:
            toc_entries = process_toc_no_page_numbers(toc_entries, page_text_by_page, toc_pages=toc_pages, model=model)
            mode = MODE_TABLE_OF_CONTENTS_WITHOUT_PAGE_NUMBERS
            steps.append("map_table_of_contents_without_page_numbers")
    else:
        toc_entries = process_no_toc(page_text_by_page, flat_md, model=model, use_llm=use_llm)
        mode = MODE_SYNTHETIC_TABLE_OF_CONTENTS
        steps.append("generate_synthetic_table_of_contents")

    toc_entries = [item for item in toc_entries if item.get("title")]
    toc_entries = _unique_by_title(toc_entries)
    if logger and logger.isEnabledFor(logging.DEBUG):
        mapped = [x for x in toc_entries if isinstance(x.get("physical_index"), int)]
        sample = [
            {"title": x.get("title"), "physical_index": x.get("physical_index")}
            for x in mapped[:8]
        ]
        logger.debug(
            "[table-of-contents] mapped_entries=%d/%d sample=%s",
            len(mapped),
            len(toc_entries),
            sample,
        )

    accuracy, incorrect = verify_toc(page_text_by_page, toc_entries, model=model)
    steps.append("verify_table_of_contents_assignments")
    if logger and logger.isEnabledFor(logging.DEBUG):
        logger.debug(
            "[table-of-contents] verification checked=%d incorrect=%d",
            len([x for x in toc_entries if isinstance(x.get("physical_index"), int)]),
            len(incorrect),
        )
        if incorrect:
            logger.debug("[table-of-contents] mismatch_sample=%s", incorrect[:5])
    if incorrect:
        toc_entries = _fix_incorrect_toc(toc_entries, incorrect, page_text_by_page, toc_pages=toc_pages)
        accuracy, incorrect = verify_toc(page_text_by_page, toc_entries, model=model)
        steps.append("repair_incorrect_table_of_contents_assignments")
        if logger and logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                "[table-of-contents] post_repair incorrect=%d accuracy=%.4f",
                len(incorrect),
                float(accuracy),
            )

    # If TOC extraction was poor, fall back to synthetic no-TOC generation.
    if toc_entries and accuracy < 0.4 and mode != MODE_SYNTHETIC_TABLE_OF_CONTENTS:
        fallback = process_no_toc(page_text_by_page, flat_md, model=model, use_llm=use_llm)
        if fallback:
            toc_entries = fallback
            mode = MODE_FALLBACK_SYNTHETIC_TABLE_OF_CONTENTS
            accuracy, incorrect = verify_toc(page_text_by_page, toc_entries, model=model)
            steps.append("fallback_to_synthetic_table_of_contents")
            if logger:
                logger.warning(
                    "[table-of-contents] low verification accuracy triggered synthetic fallback accuracy=%.4f",
                    float(accuracy),
                )

    result = OutlineResult(
        mode=mode,
        toc_pages=toc_pages,
        toc_entries=toc_entries,
        verify_accuracy=accuracy,
        verify_incorrect_count=len(incorrect),
        llm_pipeline_enabled=llm_enabled,
        steps=steps,
    )
    return result.as_dict()
