from __future__ import annotations

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

from .docling_segments import LayoutBlock


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


def _num_level(title: str) -> Optional[int]:
    m = re.match(r"^\s*(\d+(?:\.\d+)*)\b", title or "")
    if not m:
        return None
    return len(m.group(1).split("."))


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

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

    walk(nodes)
    return out


def _is_noise_heading(title: str) -> bool:
    t = (title or "").strip()
    if not t:
        return True
    if "CONTENTS*" in t.upper():
        return True
    if re.fullmatch(r"[-\s]{8,}", t):
        return True
    if re.fullmatch(r"-{3,}.+-{3,}", t):
        return True
    return False


def _is_section_like_title(title: str) -> bool:
    return _num_level(title or "") is not None


def _title_match(a: str, b: str) -> bool:
    na = _norm(a)
    nb = _norm(b)
    if not na or not nb:
        return False
    if na == nb or na in nb or nb in na:
        return True
    ta = set(na.split())
    tb = set(nb.split())
    overlap = len(ta & tb)
    threshold = max(2, int(0.7 * min(len(ta), len(tb))))
    return overlap >= threshold


def _main_anchor_titles_from_toc(toc_entries: Sequence[Dict[str, Any]]) -> List[str]:
    top_level: List[str] = []
    fallback: List[str] = []
    for item in toc_entries:
        title = str(item.get("title", "")).strip()
        if not title:
            continue
        structure = str(item.get("structure", "")).strip()
        if structure and re.fullmatch(r"\d+", structure):
            top_level.append(title)
        fallback.append(title)
    return top_level if top_level else fallback[:6]


def _split_preface_and_main(
    flat: List[Dict[str, Any]],
    toc_entries: Sequence[Dict[str, Any]],
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    if not flat:
        return [], []

    anchors = _main_anchor_titles_from_toc(toc_entries)
    if anchors:
        for i, node in enumerate(flat):
            title = str(node.get("title", "")).strip()
            if any(_title_match(title, anchor) for anchor in anchors):
                return flat[:i], flat[i:]

    # Generic fallback: first top-level numbered section (e.g., "1 ...").
    for i, node in enumerate(flat):
        if re.match(r"^\s*\d+\s+\S", str(node.get("title", ""))):
            return flat[:i], flat[i:]

    return [], flat


def _is_footer_line(text: str) -> bool:
    t = (text or "").strip()
    if not t:
        return True
    if re.match(r"^\d+\s+of\s+\d+$", t, flags=re.IGNORECASE):
        return True
    if re.match(r"^reference id\s*:\s*\d+", t, flags=re.IGNORECASE):
        return True
    return False


def _build_preface_node(
    blocks: List[LayoutBlock],
    toc_pages: Sequence[int],
    main_start_page: int,
    include_summary_fields: bool,
) -> Optional[Dict[str, Any]]:
    toc_set = {int(p) for p in toc_pages}
    preface_blocks = [b for b in blocks if b.page < main_start_page and b.page not in toc_set]
    if not preface_blocks:
        return None

    text_parts: List[str] = []
    for block in preface_blocks:
        txt = (block.text or "").strip()
        if not txt or _is_footer_line(txt):
            continue
        text_parts.append(txt)

    preface_text = "\n".join(text_parts).strip()
    if not preface_text:
        return None

    node: Dict[str, Any] = {
        "node_id": "0000",
        "title": "PREFACE / FRONT MATTER",
        "start_index": int(preface_blocks[0].page),
        "end_index": int(preface_blocks[-1].page),
        "level": 0,
        "exclusive_text": preface_text,
        "full_text": "",
        "children": [],
    }
    if include_summary_fields:
        node["summary"] = ""
        node["prefix_summary"] = None
    return node


def _semantic_rows(flat: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    stack: List[Tuple[int, str]] = []
    last_num_level: Optional[int] = None
    counter = 0

    for raw in flat:
        title = (raw.get("title") or "").strip()
        if _is_noise_heading(title):
            continue

        num_level = _num_level(title)
        if num_level is not None:
            level = num_level
            last_num_level = num_level
        elif last_num_level is not None:
            level = min(last_num_level + 1, 6)
        else:
            level = max(1, int(raw.get("level", 1) or 1))

        while stack and stack[-1][0] >= level:
            stack.pop()
        parent_id = stack[-1][1] if stack else None

        counter += 1
        node_id = str(counter).zfill(4)
        row = {
            "node_id": node_id,
            "title": title,
            "level": level,
            "parent_id": parent_id,
            "seed_text": raw.get("seed_text", ""),
        }
        rows.append(row)
        stack.append((level, node_id))

    return rows


def _rows_from_toc_entries(toc_entries: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
    entries: List[Dict[str, Any]] = []
    for item in toc_entries:
        title = str(item.get("title", "")).strip()
        if not title:
            continue
        page = item.get("physical_index")
        if not isinstance(page, int):
            continue
        structure = str(item.get("structure", "")).strip()
        level = None
        if structure and re.fullmatch(r"\d+(?:\.\d+)*", structure):
            level = len(structure.split("."))
        else:
            level = _num_level(title)
        if level is None:
            level = 1
        entries.append(
            {
                "structure": structure or None,
                "title": title,
                "level": int(level),
                "physical_index": page,
            }
        )

    if not entries:
        return []

    # Drop dangling subsection entries when parent sections are not present in the kept set.
    present_structures = {
        str(e["structure"])
        for e in entries
        if e.get("structure") and re.fullmatch(r"\d+(?:\.\d+)*", str(e["structure"]))
    }
    filtered_entries: List[Dict[str, Any]] = []
    for e in entries:
        structure = e.get("structure")
        if not structure or not re.fullmatch(r"\d+(?:\.\d+)*", str(structure)):
            filtered_entries.append(e)
            continue
        depth = len(str(structure).split("."))
        if depth <= 1:
            filtered_entries.append(e)
            continue
        parent = str(structure).rsplit(".", 1)[0]
        if parent in present_structures:
            filtered_entries.append(e)

    entries = filtered_entries
    if not entries:
        return []

    rows: List[Dict[str, Any]] = []
    counter = 0
    struct_to_id: Dict[str, str] = {}
    level_stack: List[Tuple[int, str]] = []

    for item in entries:
        counter += 1
        node_id = str(counter).zfill(4)
        level = int(item["level"])
        structure = item.get("structure")
        parent_id = None

        if structure and isinstance(structure, str) and re.fullmatch(r"\d+(?:\.\d+)*", structure):
            parent_structure = structure.rsplit(".", 1)[0] if "." in structure else None
            if parent_structure:
                parent_id = struct_to_id.get(parent_structure)
            struct_to_id[structure] = node_id
        else:
            while level_stack and level_stack[-1][0] >= level:
                level_stack.pop()
            parent_id = level_stack[-1][1] if level_stack else None

        row = {
            "node_id": node_id,
            "title": item["title"],
            "level": level,
            "parent_id": parent_id,
            "seed_text": item["title"],
        }
        rows.append(row)

        while level_stack and level_stack[-1][0] >= level:
            level_stack.pop()
        level_stack.append((level, node_id))

    return rows


def _first_index_for_page(blocks: List[LayoutBlock], page: int) -> int:
    for i, block in enumerate(blocks):
        if block.page >= page:
            return i
    return len(blocks) - 1 if blocks else 0


def _token_overlap(a: str, b: str) -> int:
    ta = set(_norm(a).split())
    tb = set(_norm(b).split())
    if not ta or not tb:
        return 0
    return len(ta & tb)


def _find_heading_index(
    title: str,
    blocks: List[LayoutBlock],
    start_idx: int,
    toc_pages: set[int],
    preferred_page: Optional[int],
) -> Tuple[int, float]:
    if not blocks:
        return 0, 0.0

    target = _norm(title)
    if not target:
        return min(start_idx, len(blocks) - 1), 0.0

    best_i = min(start_idx, len(blocks) - 1)
    best_score = -10_000.0
    target_tokens = set(target.split())

    for i in range(start_idx, len(blocks)):
        block = blocks[i]
        if block.page in toc_pages:
            continue
        txt = _norm(block.text)
        if not txt:
            continue

        score = 0.0
        if target in txt:
            score += 12.0
        if txt.startswith(target):
            score += 4.0
        overlap = len(target_tokens & set(txt.split()))
        score += float(overlap)
        if overlap < 2:
            score -= 6.0

        if preferred_page is not None:
            score -= abs(block.page - preferred_page) * 1.25

        if score > best_score:
            best_score = score
            best_i = i

        if score >= 12.0:
            break

    conf = 0.0 if best_score <= 0 else min(1.0, best_score / 14.0)
    return best_i, conf


def _build_toc_page_map(toc_entries: Sequence[Dict[str, Any]]) -> Dict[str, int]:
    out: Dict[str, int] = {}
    for item in toc_entries:
        title = _norm(str(item.get("title", "")))
        page = item.get("physical_index")
        if title and isinstance(page, int):
            out[title] = page
    return out


def _assign_spans(
    rows: List[Dict[str, Any]],
    blocks: List[LayoutBlock],
    toc_pages: Sequence[int],
    toc_entries: Sequence[Dict[str, Any]],
) -> Dict[str, Any]:
    if not rows:
        return {"heading_anchor_match_rate": 0.0}

    toc_page_set = set(int(p) for p in toc_pages)
    toc_page_map = _build_toc_page_map(toc_entries)
    content_start_page = max(toc_page_set) + 1 if toc_page_set else 1

    cursor = _first_index_for_page(blocks, content_start_page)
    while cursor < len(blocks) and blocks[cursor].page in toc_page_set:
        cursor += 1
    if cursor >= len(blocks):
        cursor = 0

    conf_sum = 0.0
    conf_count = 0

    for row in rows:
        title_norm = _norm(row["title"])
        preferred_page = toc_page_map.get(title_norm)
        search_start = cursor
        if preferred_page is not None:
            search_start = max(search_start, _first_index_for_page(blocks, preferred_page))

        idx, conf = _find_heading_index(
            row["title"],
            blocks=blocks,
            start_idx=search_start,
            toc_pages=toc_page_set,
            preferred_page=preferred_page,
        )

        row["start_i"] = max(cursor, idx)
        row["anchor_confidence"] = conf
        conf_sum += conf
        conf_count += 1
        cursor = row["start_i"]

    for i, row in enumerate(rows):
        end_i = len(blocks)
        for j in range(i + 1, len(rows)):
            if rows[j]["level"] <= row["level"]:
                end_i = rows[j]["start_i"]
                break
        if end_i <= row["start_i"]:
            end_i = min(len(blocks), row["start_i"] + 1)
        row["end_i"] = end_i

    return {
        "heading_anchor_match_rate": (conf_sum / conf_count) if conf_count else 0.0,
    }


def _build_tree(
    rows: List[Dict[str, Any]],
    blocks: List[LayoutBlock],
    include_summary_fields: bool,
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    ownership: Dict[str, List[LayoutBlock]] = {r["node_id"]: [] for r in rows}
    block_owners: Dict[str, List[str]] = {}

    for bi, block in enumerate(blocks):
        candidates = [r for r in rows if r["start_i"] <= bi < r["end_i"]]
        if not candidates:
            continue
        owner = max(candidates, key=lambda r: r["level"])
        ownership[owner["node_id"]].append(block)
        block_owners.setdefault(block.block_id, []).append(owner["node_id"])

    overlap_violations = sum(1 for owners in block_owners.values() if len(owners) > 1)
    total_block_assignments = sum(len(v) for v in block_owners.values())
    unique_blocks_assigned = len(block_owners)
    duplicated_block_rate = (
        (total_block_assignments - unique_blocks_assigned) / max(1, total_block_assignments)
        if total_block_assignments
        else 0.0
    )

    by_id: Dict[str, Dict[str, Any]] = {}
    roots: List[Dict[str, Any]] = []

    for row in rows:
        own_blocks = ownership[row["node_id"]]
        exclusive_text = "\n".join(b.text for b in own_blocks).strip()
        if not exclusive_text:
            exclusive_text = (row.get("seed_text") or "").strip()

        start_i = row["start_i"]
        end_i = max(start_i, row["end_i"] - 1)
        start_block = blocks[start_i] if blocks else None
        end_block = blocks[end_i] if blocks else None

        node: Dict[str, Any] = {
            "node_id": row["node_id"],
            "title": row["title"],
            "start_index": int(start_block.page if start_block else 1),
            "end_index": int(end_block.page if end_block else 1),
            "level": int(row["level"]),
            "exclusive_text": exclusive_text,
            "full_text": "",
            "children": [],
        }
        if include_summary_fields:
            node["summary"] = ""
            node["prefix_summary"] = None
        by_id[row["node_id"]] = node

    for row in rows:
        node = by_id[row["node_id"]]
        parent_id = row.get("parent_id")
        if parent_id and parent_id in by_id:
            by_id[parent_id]["children"].append(node)
        else:
            roots.append(node)

    return roots, {
        "overlap_violations": overlap_violations,
        "duplicated_block_rate": duplicated_block_rate,
    }


def _attach_full_text(node: Dict[str, Any]) -> str:
    child_parts = [_attach_full_text(child) for child in node.get("children", [])]
    own = (node.get("exclusive_text") or "").strip()
    parts = [p for p in [own, *child_parts] if p]
    full = "\n\n".join(parts).strip()
    node["full_text"] = full
    return full


def _reassign_node_ids(roots: List[Dict[str, Any]]) -> None:
    counter = 1

    def walk(nodes: List[Dict[str, Any]]) -> None:
        nonlocal counter
        for node in nodes:
            node["node_id"] = str(counter).zfill(4)
            counter += 1
            walk(node.get("children", []))

    walk(roots)


def _collapse_oversized_container_roots(roots: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    collapsed: List[Dict[str, Any]] = []
    for node in roots:
        children = node.get("children") or []
        if not children:
            collapsed.append(node)
            continue

        title = str(node.get("title", "")).strip()
        own_text = (node.get("exclusive_text") or "").strip()
        own_norm = _norm(own_text)
        title_norm = _norm(title)
        own_len = len(own_text)
        child_len = sum(len((c.get("exclusive_text") or "").strip()) for c in children)
        has_numbered_child = any(_is_section_like_title(str(c.get("title", ""))) for c in children)

        looks_like_header_only = own_norm == title_norm or own_len <= 120
        container_like = (not _is_section_like_title(title)) and has_numbered_child and looks_like_header_only
        tiny_vs_children = own_len <= max(120, int(0.08 * max(child_len, 1)))

        if container_like and tiny_vs_children:
            # Promote children when parent is a generic container heading with negligible own content.
            collapsed.extend(children)
        else:
            collapsed.append(node)
    return collapsed


def build_final_tree(
    markdown_tree: List[Dict[str, Any]],
    blocks: List[LayoutBlock],
    toc_pages: Sequence[int] | None = None,
    toc_entries: Sequence[Dict[str, Any]] | None = None,
    include_summary_fields: bool = True,
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    rows: List[Dict[str, Any]]
    if toc_entries:
        rows = _rows_from_toc_entries(toc_entries)
    else:
        flat = _flatten_markdown_tree(markdown_tree)
        _preface_flat, main_flat = _split_preface_and_main(flat, toc_entries=[])
        rows = _semantic_rows(main_flat)
    if not rows:
        fallback_start_page = max((int(p) for p in (toc_pages or [])), default=0) + 1
        preface_only = _build_preface_node(
            blocks,
            toc_pages=toc_pages or [],
            main_start_page=fallback_start_page,
            include_summary_fields=include_summary_fields,
        )
        if preface_only is None:
            return [], {"heading_anchor_match_rate": 0.0, "overlap_violations": 0, "duplicated_block_rate": 0.0}
        _attach_full_text(preface_only)
        roots = [preface_only]
        _reassign_node_ids(roots)
        return roots, {
            "heading_anchor_match_rate": 0.0,
            "overlap_violations": 0,
            "duplicated_block_rate": 0.0,
            "preface_included": True,
            "preface_pages": preface_only["end_index"] - preface_only["start_index"] + 1,
        }

    assign_meta = _assign_spans(
        rows,
        blocks,
        toc_pages=toc_pages or [],
        toc_entries=toc_entries or [],
    )
    roots, build_meta = _build_tree(rows, blocks, include_summary_fields=include_summary_fields)
    main_start_page = min((r.get("start_index", 1) for r in roots), default=1)
    preface_node = _build_preface_node(
        blocks,
        toc_pages=toc_pages or [],
        main_start_page=main_start_page,
        include_summary_fields=include_summary_fields,
    )
    if preface_node is not None:
        roots = [preface_node, *roots]

    roots = _collapse_oversized_container_roots(roots)
    roots.sort(key=lambda n: (int(n.get("start_index", 0)), int(n.get("level", 0))))
    for root in roots:
        _attach_full_text(root)
    _reassign_node_ids(roots)

    return roots, {
        **assign_meta,
        **build_meta,
        "preface_included": preface_node is not None,
        "preface_pages": (preface_node["end_index"] - preface_node["start_index"] + 1) if preface_node else 0,
    }


def validation_metrics(structure: List[Dict[str, Any]]) -> Dict[str, Any]:
    node_count = 0
    empty_count = 0

    def walk(nodes: List[Dict[str, Any]]) -> None:
        nonlocal node_count, empty_count
        for n in nodes:
            node_count += 1
            if not (n.get("full_text") or "").strip():
                empty_count += 1
            walk(n.get("children", []))

    walk(structure)
    return {
        "node_count": node_count,
        "empty_node_rate": (empty_count / node_count) if node_count else 0.0,
    }
