from __future__ import annotations

import math
from typing import Any, Dict, List, Sequence, Tuple

from ..common.utils import count_tokens


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

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

    walk(nodes)


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 _split_by_paragraphs(text: str) -> List[str]:
    parts = [p.strip() for p in text.split("\n\n") if p.strip()]
    if parts:
        return parts
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    if lines:
        return lines
    return [text.strip()] if text.strip() else []


def _chunk_parts(parts: Sequence[str], max_tokens: int) -> List[str]:
    def split_long_part(part: str) -> List[str]:
        words = part.split()
        if not words:
            return []
        slices: List[str] = []
        buf: List[str] = []
        for word in words:
            candidate = " ".join([*buf, word]).strip()
            if buf and count_tokens(candidate, model="gpt-4o") > max_tokens:
                slices.append(" ".join(buf).strip())
                buf = [word]
            else:
                buf.append(word)
        if buf:
            slices.append(" ".join(buf).strip())
        return [s for s in slices if s]

    chunks: List[str] = []
    buf: List[str] = []
    buf_tokens = 0
    for part in parts:
        tok = count_tokens(part, model="gpt-4o")
        if tok > max_tokens:
            small_parts = split_long_part(part)
            for small in small_parts:
                small_tok = count_tokens(small, model="gpt-4o")
                if buf and buf_tokens + small_tok > max_tokens:
                    chunks.append("\n\n".join(buf).strip())
                    buf = [small]
                    buf_tokens = small_tok
                else:
                    buf.append(small)
                    buf_tokens += small_tok
            continue

        if buf and buf_tokens + tok > max_tokens:
            chunks.append("\n\n".join(buf).strip())
            buf = [part]
            buf_tokens = tok
        else:
            buf.append(part)
            buf_tokens += tok
    if buf:
        chunks.append("\n\n".join(buf).strip())
    return [chunk for chunk in chunks if chunk]


def _split_pages(start: int, end: int, chunk_count: int) -> List[Tuple[int, int]]:
    if chunk_count <= 1:
        return [(start, end)]
    span = max(1, end - start + 1)
    per = max(1, math.ceil(span / chunk_count))
    ranges: List[Tuple[int, int]] = []
    current = start
    while current <= end and len(ranges) < chunk_count:
        r_end = min(end, current + per - 1)
        ranges.append((current, r_end))
        current = r_end + 1
    while len(ranges) < chunk_count:
        ranges.append((end, end))
    return ranges[:chunk_count]


def _make_part_node(
    *,
    parent: Dict[str, Any],
    text: str,
    index: int,
    page_range: Tuple[int, int],
    include_summary_fields: bool,
) -> Dict[str, Any]:
    part_node: Dict[str, Any] = {
        "node_id": "",
        "title": f"{parent.get('title', 'Section')} (Part {index})",
        "start_index": int(page_range[0]),
        "end_index": int(page_range[1]),
        "level": int(parent.get("level", 1)) + 1,
        "exclusive_text": text.strip(),
        "full_text": "",
        "children": [],
    }
    if include_summary_fields:
        part_node["summary"] = ""
        part_node["prefix_summary"] = None
    return part_node


def _maybe_split_node(
    node: Dict[str, Any],
    *,
    max_token_num_each_node: int,
    max_page_num_each_node: int,
    include_summary_fields: bool,
) -> Dict[str, int]:
    stats = {"nodes_split": 0, "parts_created": 0}

    own_text = (node.get("exclusive_text") or "").strip()
    if own_text:
        token_count = count_tokens(own_text, model="gpt-4o")
        start = int(node.get("start_index") or 1)
        end = int(node.get("end_index") or start)
        page_count = max(1, end - start + 1)

        token_parts = math.ceil(token_count / max(1, max_token_num_each_node))
        page_parts = math.ceil(page_count / max(1, max_page_num_each_node))
        needed_parts = max(token_parts, page_parts)

        if needed_parts > 1:
            parts = _split_by_paragraphs(own_text)
            chunked = _chunk_parts(parts, max_tokens=max(1, max_token_num_each_node))
            if len(chunked) < needed_parts:
                needed_parts = len(chunked)

            if needed_parts > 1 and len(chunked) > 1:
                page_ranges = _split_pages(start, end, len(chunked))
                split_children = [
                    _make_part_node(
                        parent=node,
                        text=chunk,
                        index=i + 1,
                        page_range=page_ranges[i],
                        include_summary_fields=include_summary_fields,
                    )
                    for i, chunk in enumerate(chunked)
                ]
                existing_children = node.get("children", []) if isinstance(node.get("children"), list) else []
                node["exclusive_text"] = ""
                node["children"] = [*split_children, *existing_children]
                stats["nodes_split"] += 1
                stats["parts_created"] += len(split_children)

    children = node.get("children", [])
    if isinstance(children, list):
        for child in children:
            if not isinstance(child, dict):
                continue
            child_stats = _maybe_split_node(
                child,
                max_token_num_each_node=max_token_num_each_node,
                max_page_num_each_node=max_page_num_each_node,
                include_summary_fields=include_summary_fields,
            )
            stats["nodes_split"] += child_stats["nodes_split"]
            stats["parts_created"] += child_stats["parts_created"]

    return stats


def enforce_node_limits(
    structure: List[Dict[str, Any]],
    *,
    max_token_num_each_node: int,
    max_page_num_each_node: int,
    include_summary_fields: bool,
) -> Dict[str, int]:
    if not structure:
        return {"nodes_split": 0, "parts_created": 0}

    stats = {"nodes_split": 0, "parts_created": 0}
    for node in structure:
        if not isinstance(node, dict):
            continue
        child_stats = _maybe_split_node(
            node,
            max_token_num_each_node=max_token_num_each_node,
            max_page_num_each_node=max_page_num_each_node,
            include_summary_fields=include_summary_fields,
        )
        stats["nodes_split"] += child_stats["nodes_split"]
        stats["parts_created"] += child_stats["parts_created"]

    for node in structure:
        _attach_full_text(node)
    _reassign_node_ids(structure)
    return stats


def build_structure_from_sections(
    sections: List[Dict[str, Any]],
    *,
    include_summary_fields: bool,
) -> List[Dict[str, Any]]:
    def build_node(section: Dict[str, Any], default_level: int = 1) -> Dict[str, Any]:
        level = int(section.get("level") or default_level)
        node: Dict[str, Any] = {
            "node_id": "",
            "title": str(section.get("title", "Section")).strip() or "Section",
            "start_index": int(section.get("start_index") or 1),
            "end_index": int(section.get("end_index") or int(section.get("start_index") or 1)),
            "level": level,
            "exclusive_text": str(section.get("text", "")).strip(),
            "full_text": "",
            "children": [],
        }
        if include_summary_fields:
            node["summary"] = ""
            node["prefix_summary"] = None
        if section.get("start_char") is not None:
            node["start_char"] = int(section["start_char"])
        if section.get("end_char") is not None:
            node["end_char"] = int(section["end_char"])
        topic_meta = section.get("topic_meta")
        if isinstance(topic_meta, dict) and topic_meta:
            node["topic_meta"] = topic_meta

        children = section.get("children")
        if isinstance(children, list):
            for child in children:
                if not isinstance(child, dict):
                    continue
                child_level = int(child.get("level") or (level + 1))
                node["children"].append(build_node(child, default_level=child_level))
        return node

    nodes: List[Dict[str, Any]] = []
    for idx, section in enumerate(sections, start=1):
        if not isinstance(section, dict):
            continue
        built = build_node(section, default_level=int(section.get("level") or 1))
        if not built.get("title"):
            built["title"] = f"Section {idx}"
        nodes.append(built)

    for node in nodes:
        _attach_full_text(node)
    _reassign_node_ids(nodes)
    return nodes
