from __future__ import annotations

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

from ..common.utils import ChatGPT_API, count_tokens


def _title_from_text(text: str, fallback: str) -> str:
    lines = [ln.strip() for ln in (text or "").splitlines() if ln.strip()]
    for line in lines[:12]:
        if 3 <= len(line.split()) <= 10 and len(line) <= 80:
            return line
    return fallback


def _chunk_pages(
    page_text_by_page: Dict[int, str],
    max_tokens: int,
    max_pages: int,
) -> List[Tuple[int, int, str]]:
    pages = sorted(page_text_by_page)
    if not pages:
        return []

    chunks: List[Tuple[int, int, str]] = []
    buf_pages: List[int] = []
    buf_parts: List[str] = []
    buf_tokens = 0

    for page in pages:
        page_text = (page_text_by_page.get(page) or "").strip()
        wrapped = f"[Page {page}]\n{page_text}"
        wrapped_tokens = count_tokens(wrapped, model="gpt-4o")

        exceeds_pages = len(buf_pages) >= max_pages
        exceeds_tokens = (buf_tokens + wrapped_tokens) > max_tokens if buf_pages else False
        if buf_pages and (exceeds_pages or exceeds_tokens):
            chunks.append((buf_pages[0], buf_pages[-1], "\n\n".join(buf_parts).strip()))
            buf_pages = []
            buf_parts = []
            buf_tokens = 0

        buf_pages.append(page)
        buf_parts.append(wrapped)
        buf_tokens += wrapped_tokens

    if buf_pages:
        chunks.append((buf_pages[0], buf_pages[-1], "\n\n".join(buf_parts).strip()))
    return chunks


def _llm_heading(
    text: str,
    model: Optional[str],
) -> Optional[str]:
    if not model:
        return None
    prompt = (
        "Create one concise heading (3-8 words) for this document section. "
        "Return only the heading text.\n\n"
        f"Section text:\n{text}"
    )
    raw = ChatGPT_API(model=model, prompt=prompt)
    if raw == "Error":
        return None
    heading = re.sub(r"\s+", " ", str(raw).strip().strip('"').strip("'"))
    return heading or None


def build_unstructured_sections(
    *,
    page_text_by_page: Dict[int, str],
    model: Optional[str],
    use_llm: bool,
    max_token_num_each_node: int,
    max_page_num_each_node: int,
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    chunks = _chunk_pages(
        page_text_by_page=page_text_by_page,
        max_tokens=max(512, max_token_num_each_node),
        max_pages=max(1, max_page_num_each_node),
    )

    sections: List[Dict[str, Any]] = []
    for idx, (start_page, end_page, chunk_text) in enumerate(chunks, start=1):
        fallback = f"Section {idx}"
        heading = _llm_heading(chunk_text[:8000], model=model) if use_llm else None
        title = heading or _title_from_text(chunk_text, fallback=fallback)
        sections.append(
            {
                "title": title,
                "text": chunk_text,
                "start_index": start_page,
                "end_index": end_page,
                "level": 1,
            }
        )

    meta = {
        "mode": "unstructured_synthetic_sections",
        "table_of_contents_pages": [],
        "table_of_contents_entries": [
            {
                "structure": str(i),
                "title": row["title"],
                "page": None,
                "physical_index": row["start_index"],
            }
            for i, row in enumerate(sections, start=1)
        ],
        "verify_accuracy": 1.0 if sections else 0.0,
        "verify_incorrect_count": 0,
        "llm_pipeline_enabled": bool(use_llm),
        "steps": ["generate_unstructured_sections"],
        "generated_section_count": len(sections),
    }
    return sections, meta
