from __future__ import annotations

import math
import re
from typing import Any, Dict, List, 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 _token_count(text: str) -> int:
    return count_tokens(text, model="gpt-4o")


def _split_plain_text(text: str, max_tokens: int) -> List[str]:
    parts = [p.strip() for p in re.split(r"(?<=[.!?])\s+|\n{2,}", text) if p and p.strip()]
    if not parts:
        parts = [ln.strip() for ln in text.splitlines() if ln.strip()]
    if not parts:
        parts = [text.strip()] if text.strip() else []

    chunks: List[str] = []
    buf: List[str] = []
    buf_tokens = 0
    for part in parts:
        part_tokens = _token_count(part)
        if part_tokens > max_tokens:
            words = part.split()
            word_buf: List[str] = []
            for word in words:
                candidate = " ".join([*word_buf, word]).strip()
                if word_buf and _token_count(candidate) > max_tokens:
                    chunks.append(" ".join(word_buf).strip())
                    word_buf = [word]
                else:
                    word_buf.append(word)
            if word_buf:
                part = " ".join(word_buf).strip()
                part_tokens = _token_count(part)
            else:
                continue

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


def _split_markdown_blocks(raw_text: str) -> List[str]:
    if not raw_text.strip():
        return []
    lines = raw_text.splitlines(keepends=True)
    blocks: List[str] = []
    buf: List[str] = []
    in_fence = False
    fence_delim = ""

    for line in lines:
        match = re.match(r"^\s*(`{3,}|~{3,})", line)
        if match:
            delim = match.group(1)
            if not in_fence:
                in_fence = True
                fence_delim = delim[0] * len(delim)
                buf.append(line)
                continue
            if line.lstrip().startswith(fence_delim):
                buf.append(line)
                in_fence = False
                fence_delim = ""
                continue

        if in_fence:
            buf.append(line)
            continue

        if line.strip() == "":
            if buf:
                blocks.append("".join(buf).rstrip("\n"))
                buf = []
            continue

        buf.append(line)

    if buf:
        blocks.append("".join(buf).rstrip("\n"))
    return [b for b in blocks if b.strip()]


def _is_heading_block(block: str) -> bool:
    first = block.splitlines()[0].strip() if block.splitlines() else ""
    return bool(re.match(r"^#{1,6}\s+\S+", first))


def _is_list_block(block: str) -> bool:
    lines = [ln.strip() for ln in block.splitlines() if ln.strip()]
    if not lines:
        return False
    return all(bool(re.match(r"^([-*+]|\d+[.)])\s+", ln)) for ln in lines)


def _is_table_block(block: str) -> bool:
    lines = [ln.rstrip() for ln in block.splitlines() if ln.strip()]
    if len(lines) < 2:
        return False
    if "|" not in lines[0] or "|" not in lines[1]:
        return False
    return bool(re.match(r"^\s*[:\-| ]+\s*$", lines[1]))


def _is_code_block(block: str) -> bool:
    first = block.splitlines()[0].lstrip() if block.splitlines() else ""
    return first.startswith("```") or first.startswith("~~~")


def _is_image_block(block: str) -> bool:
    first = block.splitlines()[0].strip() if block.splitlines() else ""
    if re.match(r"^!\[.*\]\(.*\)$", first):
        return True
    return first.lower() == "<!-- image -->"


def _classify_block(block: str) -> str:
    if _is_heading_block(block):
        return "heading"
    if _is_list_block(block):
        return "list"
    if _is_table_block(block):
        return "table"
    if _is_code_block(block):
        return "code"
    if _is_image_block(block):
        return "image"
    return "paragraph"


def _split_list_unit(text: str, max_tokens: int) -> List[str]:
    lines = [ln for ln in text.splitlines() if ln.strip()]
    chunks: List[str] = []
    buf: List[str] = []
    buf_tokens = 0
    for line in lines:
        line_tokens = _token_count(line)
        if line_tokens > max_tokens:
            chunks.extend(_split_plain_text(line, max_tokens))
            continue
        if buf and (buf_tokens + line_tokens > max_tokens):
            chunks.append("\n".join(buf).strip())
            buf = [line]
            buf_tokens = line_tokens
        else:
            buf.append(line)
            buf_tokens += line_tokens
    if buf:
        chunks.append("\n".join(buf).strip())
    return [c for c in chunks if c]


def _split_table_unit(text: str, max_tokens: int) -> List[str]:
    lines = [ln.rstrip() for ln in text.splitlines() if ln.strip()]
    if len(lines) < 3:
        return _split_plain_text(text, max_tokens)
    header = lines[0]
    sep = lines[1]
    rows = lines[2:]
    header_tokens = _token_count(f"{header}\n{sep}")
    if header_tokens >= max_tokens:
        return _split_plain_text(text, max_tokens)

    chunks: List[str] = []
    cur_rows: List[str] = []
    cur_tokens = header_tokens
    for row in rows:
        row_tokens = _token_count(row)
        if row_tokens > max_tokens:
            if cur_rows:
                chunks.append("\n".join([header, sep, *cur_rows]).strip())
                cur_rows = []
                cur_tokens = header_tokens
            for small in _split_plain_text(row, max_tokens - header_tokens):
                chunks.append("\n".join([header, sep, small]).strip())
            continue

        if cur_rows and (cur_tokens + row_tokens > max_tokens):
            chunks.append("\n".join([header, sep, *cur_rows]).strip())
            cur_rows = [row]
            cur_tokens = header_tokens + row_tokens
        else:
            cur_rows.append(row)
            cur_tokens += row_tokens

    if cur_rows:
        chunks.append("\n".join([header, sep, *cur_rows]).strip())
    return [c for c in chunks if c]


def _split_code_unit(text: str, max_tokens: int) -> List[str]:
    lines = text.splitlines()
    if not lines:
        return []
    opener = lines[0]
    closer = lines[-1] if len(lines) > 1 and lines[-1].lstrip().startswith(opener.strip()[:3]) else opener[:3]
    body = lines[1:-1] if len(lines) > 2 else lines[1:]
    frame_tokens = _token_count(f"{opener}\n{closer}")
    if frame_tokens >= max_tokens:
        return _split_plain_text(text, max_tokens)

    chunks: List[str] = []
    buf: List[str] = []
    buf_tokens = 0
    for line in body:
        line_tokens = _token_count(line)
        if line_tokens > max_tokens - frame_tokens:
            if buf:
                chunks.append("\n".join([opener, *buf, closer]).strip())
                buf = []
                buf_tokens = 0
            for small in _split_plain_text(line, max_tokens - frame_tokens):
                chunks.append("\n".join([opener, small, closer]).strip())
            continue

        if buf and (buf_tokens + line_tokens + frame_tokens > max_tokens):
            chunks.append("\n".join([opener, *buf, closer]).strip())
            buf = [line]
            buf_tokens = line_tokens
        else:
            buf.append(line)
            buf_tokens += line_tokens

    if buf:
        chunks.append("\n".join([opener, *buf, closer]).strip())
    return [c for c in chunks if c]


def _split_unit(kind: str, text: str, max_tokens: int) -> List[str]:
    if kind == "list":
        return _split_list_unit(text, max_tokens)
    if kind == "table":
        return _split_table_unit(text, max_tokens)
    if kind == "code":
        return _split_code_unit(text, max_tokens)
    return _split_plain_text(text, max_tokens)


def _chunk_markdown_text(text: str, max_tokens: int) -> List[str]:
    blocks = _split_markdown_blocks(text)
    if not blocks:
        return []
    units = [(_classify_block(block), block) for block in blocks]

    chunks: List[str] = []
    cur_parts: List[str] = []
    cur_tokens = 0

    idx = 0
    while idx < len(units):
        kind, body = units[idx]
        body_tokens = _token_count(body)

        # Keep headings with following content whenever possible.
        if kind == "heading":
            if cur_parts:
                chunks.append("\n\n".join(cur_parts).strip())
                cur_parts = []
                cur_tokens = 0
            if body_tokens > max_tokens:
                chunks.extend(_split_unit("paragraph", body, max_tokens))
                idx += 1
                continue
            cur_parts = [body]
            cur_tokens = body_tokens
            if idx + 1 < len(units):
                next_kind, next_body = units[idx + 1]
                next_tokens = _token_count(next_body)
                if next_kind != "heading" and next_tokens <= max_tokens and cur_tokens + next_tokens <= max_tokens:
                    cur_parts.append(next_body)
                    cur_tokens += next_tokens
                    idx += 1
            idx += 1
            continue

        if body_tokens <= max_tokens:
            if cur_parts and cur_tokens + body_tokens > max_tokens:
                chunks.append("\n\n".join(cur_parts).strip())
                cur_parts = [body]
                cur_tokens = body_tokens
            else:
                cur_parts.append(body)
                cur_tokens += body_tokens
            idx += 1
            continue

        if cur_parts:
            chunks.append("\n\n".join(cur_parts).strip())
            cur_parts = []
            cur_tokens = 0

        subchunks = _split_unit(kind, body, max_tokens)
        for sub in subchunks:
            sub_tokens = _token_count(sub)
            if sub_tokens > max_tokens:
                for fallback in _split_plain_text(sub, max_tokens):
                    if fallback.strip():
                        chunks.append(fallback.strip())
                continue
            if cur_parts and cur_tokens + sub_tokens > max_tokens:
                chunks.append("\n\n".join(cur_parts).strip())
                cur_parts = [sub]
                cur_tokens = sub_tokens
            else:
                cur_parts.append(sub)
                cur_tokens += sub_tokens
        idx += 1

    if cur_parts:
        chunks.append("\n\n".join(cur_parts).strip())

    return [chunk for chunk in chunks if chunk.strip()]


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:
            chunked = _chunk_markdown_text(own_text, 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
