from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import logging
import os
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from ..adapters.legacy_compat import navexa_to_legacy_compat
from ..common.env import load_navexa_env
from ..common.utils import generate_node_summary, get_usage_summary
from ..llm.client import LLMSettings, has_llm_credentials, reset_llm_client, resolve_model_name
from ..outline.markdown_outline import markdown_to_tree_from_content
from ..outline.outline_engine import build_outline_engine
from ..outline.semi_structured_outline import build_semi_structured_outline
from ..outline.transcript_outline import build_transcript_sections
from ..outline.unstructured_outline import build_unstructured_sections
from ..segment.docling_segments import extract_docling_content
from ..segment.node_limits import build_structure_from_sections, enforce_node_limits
from ..segment.span_assign import build_final_tree, validation_metrics
from ..segment.transcript_text import extract_transcript_text_content

PIPELINE_VERSION = "navexa-1.2.0"
MODE_LLM = "llm"
MODE_NO_LLM = "no-llm"
VERBOSITY_LOW = "low"
VERBOSITY_MEDIUM = "medium"
VERBOSITY_HIGH = "high"
DOCUMENT_TYPE_STRUCTURED = "structured"
DOCUMENT_TYPE_SEMI_STRUCTURED = "semi_structured"
DOCUMENT_TYPE_UNSTRUCTURED = "unstructured"
DOCUMENT_TYPE_TRANSCRIPT = "transcript"

load_navexa_env()


def _normalize_mode(mode: Optional[str]) -> str:
    raw = (mode or os.getenv("NAVEXA_MODE", MODE_NO_LLM)).strip().lower()
    if raw in {"llm", "use-llm", "with-llm"}:
        return MODE_LLM
    return MODE_NO_LLM


def _normalize_verbosity(value: Optional[str]) -> str:
    raw = (value or os.getenv("NAVEXA_VERBOSE", VERBOSITY_MEDIUM)).strip().lower()
    if raw in {"1", "low"}:
        return VERBOSITY_LOW
    if raw in {"3", "high", "debug", "detailed"}:
        return VERBOSITY_HIGH
    return VERBOSITY_MEDIUM


def _normalize_document_type(value: Optional[str]) -> str:
    raw = (value or os.getenv("NAVEXA_DOCUMENT_TYPE", DOCUMENT_TYPE_STRUCTURED)).strip().lower()
    if raw in {"semi", "semi-structured", "semistructured"}:
        return DOCUMENT_TYPE_SEMI_STRUCTURED
    if raw in {"unstructured", "un-structured"}:
        return DOCUMENT_TYPE_UNSTRUCTURED
    if raw in {"transcript", "transcripts"}:
        return DOCUMENT_TYPE_TRANSCRIPT
    return DOCUMENT_TYPE_STRUCTURED


def _normalize_summary_toggle(value: Optional[str]) -> bool:
    raw = (value or os.getenv("NAVEXA_IF_ADD_NODE_SUMMARY", "yes")).strip().lower()
    return raw in {"1", "true", "yes", "on", "y"}


def _verbosity_to_level(verbosity: str) -> int:
    if verbosity == VERBOSITY_LOW:
        return logging.INFO
    if verbosity == VERBOSITY_HIGH:
        return logging.DEBUG
    return logging.INFO


def _configure_logger(verbosity: str) -> logging.Logger:
    logger = logging.getLogger("navexa")
    level = _verbosity_to_level(verbosity)
    logger.setLevel(level)
    logger.propagate = False

    if not logger.handlers:
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter("[NAVEXA][%(levelname)s] %(message)s"))
        logger.addHandler(handler)

    for handler in logger.handlers:
        handler.setLevel(level)
    return logger


def _log_banner(logger: logging.Logger, title: str) -> None:
    line = "=" * 72
    logger.info(line)
    logger.info("[ %s ]", title)
    logger.info(line)


def _summarize_top_level_titles(structure: List[Dict[str, Any]], limit: int = 8) -> List[str]:
    titles = [str(node.get("title", "")).strip() for node in structure]
    return [t for t in titles if t][:limit]


def _has_llm_credentials() -> bool:
    return has_llm_credentials()


def _resolve_model(user_model: Optional[str], use_llm_requested: bool) -> Optional[str]:
    return resolve_model_name(user_model, use_llm_requested=use_llm_requested)


def _sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _slug(path: str) -> str:
    stem = Path(path).stem.lower()
    return "".join(ch if ch.isalnum() else "-" for ch in stem).strip("-") or "document"


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

    def walk(items: List[Dict[str, Any]]) -> None:
        for node in items:
            out.append(node)
            walk(node.get("children", []))

    walk(nodes)
    return out


def _drop_summary_fields(nodes: List[Dict[str, Any]]) -> None:
    for node in _flatten(nodes):
        node.pop("summary", None)
        node.pop("prefix_summary", None)


def _run_async(coro: Any) -> Any:
    """Run coroutine in both script and notebook environments."""
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(coro)

    result_holder: Dict[str, Any] = {}
    error_holder: Dict[str, BaseException] = {}

    def worker() -> None:
        try:
            result_holder["value"] = asyncio.run(coro)
        except BaseException as exc:  # propagate original exception across thread boundary
            error_holder["error"] = exc

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()
    thread.join()

    if "error" in error_holder:
        raise error_holder["error"]
    return result_holder.get("value")


def _clear_summaries(structure: List[Dict[str, Any]]) -> None:
    for node in _flatten(structure):
        node["summary"] = ""
        node["prefix_summary"] = None


async def _summarize(
    structure: List[Dict[str, Any]],
    model: Optional[str],
    use_llm: bool,
    logger: logging.Logger,
    include_summary_fields: bool,
) -> None:
    nodes = _flatten(structure)
    if not nodes:
        return

    if not include_summary_fields:
        _drop_summary_fields(structure)
        logger.info("[summary] disabled by parameter if_add_node_summary=no")
        return

    summary_disabled = os.getenv("NAVEXA_DISABLE_SUMMARY", "0").strip().lower() in {"1", "true", "yes", "on"}
    if not use_llm or summary_disabled:
        _clear_summaries(structure)
        if summary_disabled:
            logger.info("[summary] disabled by NAVEXA_DISABLE_SUMMARY")
        else:
            logger.info("[summary] mode=no-llm -> summaries left empty")
        return

    if not model or not _has_llm_credentials():
        _clear_summaries(structure)
        logger.warning("[summary] LLM summary requested but model/key is missing. Summaries left empty.")
        return

    # Cost optimization: when parent has exactly one child, reuse child summary
    # for the parent and skip a separate parent LLM summary call.
    copy_from_single_child: Dict[str, str] = {}
    for node in nodes:
        children = node.get("children", [])
        if not isinstance(children, list) or len(children) != 1:
            continue
        child_id = (children[0] or {}).get("node_id") if isinstance(children[0], dict) else None
        if not child_id:
            continue
        copy_from_single_child[str(node.get("node_id"))] = str(child_id)

    summarize_nodes = [n for n in nodes if str(n.get("node_id")) not in copy_from_single_child]
    logger.info(
        "[summary] generating LLM summaries for %d nodes (copied_from_child=%d)",
        len(summarize_nodes),
        len(copy_from_single_child),
    )
    semaphore = asyncio.Semaphore(4)

    async def one(node: Dict[str, Any]) -> Tuple[str, str, Optional[str]]:
        inclusive = (node.get("full_text") or "").strip()
        exclusive = (node.get("exclusive_text") or "").strip()
        has_children = bool(node.get("children"))
        if not inclusive and not exclusive:
            return node["node_id"], "", None

        async with semaphore:
            summary = await generate_node_summary({"full_text": inclusive}, model=model)
        if summary == "Error":
            summary = ""

        prefix_summary: Optional[str] = None
        if has_children and exclusive:
            async with semaphore:
                prefix_summary = await generate_node_summary({"full_text": exclusive}, model=model)
            if prefix_summary == "Error":
                prefix_summary = None

        return node["node_id"], summary, prefix_summary

    results = await asyncio.gather(*[one(node) for node in summarize_nodes])
    by_id = {node["node_id"]: node for node in nodes}
    for node_id, summary, prefix_summary in results:
        by_id[node_id]["summary"] = summary
        by_id[node_id]["prefix_summary"] = prefix_summary

    def _apply_copy(node: Dict[str, Any]) -> None:
        for child in node.get("children", []):
            if isinstance(child, dict):
                _apply_copy(child)
        node_id = str(node.get("node_id"))
        child_id = copy_from_single_child.get(node_id)
        if not child_id:
            return
        child = by_id.get(child_id)
        if not child:
            return
        node["summary"] = child.get("summary", "")
        node["prefix_summary"] = None

    for root in structure:
        if isinstance(root, dict):
            _apply_copy(root)


def build_navexa_document(
    pdf_path: str,
    model: Optional[str] = None,
    mode: Optional[str] = None,
    verbosity: Optional[str] = None,
    parser_model: Optional[str] = None,
    output_format: Optional[str] = None,
    document_type: Optional[str] = None,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: Optional[str] = "yes",
    semi_heading_prompt_template: Optional[str] = None,
    transcript_topic_prompt_template: Optional[str] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
    resolved_mode = _normalize_mode(mode)
    resolved_verbosity = _normalize_verbosity(verbosity)
    resolved_document_type = _normalize_document_type(document_type)
    summary_enabled = _normalize_summary_toggle(if_add_node_summary)
    logger = _configure_logger(resolved_verbosity)
    show_steps = resolved_verbosity in {VERBOSITY_MEDIUM, VERBOSITY_HIGH}
    show_details = resolved_verbosity == VERBOSITY_HIGH

    _log_banner(logger, "Navexa Document Indexing")
    logger.info("[config] pdf=%s", pdf_path)
    logger.info("[config] mode=%s verbosity=%s", resolved_mode, resolved_verbosity)
    logger.info("[config] document_type=%s", resolved_document_type)
    logger.info("[config] if_add_node_summary=%s", "yes" if summary_enabled else "no")
    resolved_parser_model = (parser_model or os.getenv("NAVEXA_PARSER_MODEL", "docling")).strip().lower() or "docling"
    resolved_output_format = (output_format or os.getenv("NAVEXA_DOCLING_OUTPUT_FORMAT", "markdown")).strip().lower() or "markdown"
    logger.info("[config] parser_model=%s output_format=%s", resolved_parser_model, resolved_output_format)

    llm_required = resolved_document_type in {DOCUMENT_TYPE_SEMI_STRUCTURED, DOCUMENT_TYPE_TRANSCRIPT}
    use_llm_requested = resolved_mode == MODE_LLM or llm_required
    selected_model = _resolve_model(model, use_llm_requested=use_llm_requested)
    has_credentials = _has_llm_credentials()
    use_llm_effective = use_llm_requested and bool(selected_model) and has_credentials

    if resolved_mode == MODE_LLM and not has_credentials:
        raise RuntimeError(
            "mode='llm' requested but credentials are missing for selected provider. "
            "Use OPENAI_API_KEY (openai) or AZURE_OPENAI_API_KEY + AZURE_OPENAI_BASE_URL (azure)."
        )
    if resolved_mode == MODE_LLM and not selected_model:
        raise RuntimeError(
            "mode='llm' requested but model is not configured. "
            "Pass model=... explicitly, or set OPENAI_MODEL_NAME (openai) / AZURE_DEPLOYMENT_NAME (azure)."
        )
    if llm_required and not use_llm_effective:
        raise RuntimeError(
            f"document_type='{resolved_document_type}' requires LLM, but model/API key is missing."
        )

    logger.info(
        "[config] effective_mode=%s model=%s",
        MODE_LLM if use_llm_effective else MODE_NO_LLM,
        selected_model or "none",
    )

    reset_llm_client(LLMSettings(model=selected_model))

    markdown = ""
    blocks: List[Any] = []
    if resolved_document_type == DOCUMENT_TYPE_TRANSCRIPT:
        if show_steps:
            logger.info("[step] transcript text extraction")
        page_text_by_page, total_pages = extract_transcript_text_content(pdf_path)
        if show_details:
            logger.debug("[transcript] pages=%d", total_pages)
            logger.debug("[transcript] extracted_page_text_entries=%d", len(page_text_by_page))
    else:
        if show_steps:
            logger.info("[step] docling extraction")
        markdown, blocks, page_text_by_page, total_pages = extract_docling_content(
            pdf_path,
            output_format=resolved_output_format,
            parser_model=resolved_parser_model,
        )
        if show_details:
            logger.debug("[docling] pages=%d blocks=%d", total_pages, len(blocks))
            logger.debug("[docling] extracted_page_text_entries=%d", len(page_text_by_page))

    md_tree: Dict[str, Any] = {"structure": []}
    if resolved_document_type in {DOCUMENT_TYPE_STRUCTURED, DOCUMENT_TYPE_SEMI_STRUCTURED}:
        if show_steps:
            logger.info("[step] markdown heading tree")
        md_tree = _run_async(markdown_to_tree_from_content(markdown))
        if show_details:
            logger.debug("[markdown] root_nodes=%d", len(md_tree.get("structure", []) or []))

    outline_meta: Dict[str, Any]
    segment_meta: Dict[str, Any]
    structure: List[Dict[str, Any]]

    if resolved_document_type in {DOCUMENT_TYPE_STRUCTURED, DOCUMENT_TYPE_SEMI_STRUCTURED}:
        if show_steps:
            logger.info("[step] table of contents pipeline")
        base_outline_meta = build_outline_engine(
            markdown_tree=md_tree.get("structure", []),
            page_text_by_page=page_text_by_page,
            model=selected_model,
            use_llm=use_llm_effective,
            logger=logger if show_steps else None,
        )
        if resolved_document_type == DOCUMENT_TYPE_STRUCTURED:
            outline_meta = base_outline_meta
        else:
            outline_meta = build_semi_structured_outline(
                base_outline=base_outline_meta,
                page_text_by_page=page_text_by_page,
                model=selected_model,
                use_llm=use_llm_effective,
                prompt_template=semi_heading_prompt_template,
            )

        if show_steps:
            logger.info(
                "[table-of-contents] mode=%s pages=%s accuracy=%.2f",
                outline_meta.get("mode"),
                outline_meta.get("table_of_contents_pages", []),
                float(outline_meta.get("verify_accuracy", 0.0)),
            )
        if show_details:
            logger.debug(
                "[table-of-contents] entries=%d",
                len(outline_meta.get("table_of_contents_entries", []) or []),
            )

        if show_steps:
            logger.info("[step] span segmentation")
        structure, segment_meta = build_final_tree(
            markdown_tree=md_tree.get("structure", []),
            blocks=blocks,
            toc_pages=outline_meta.get("table_of_contents_pages", []),
            toc_entries=outline_meta.get("table_of_contents_entries", []),
            include_summary_fields=summary_enabled,
        )
    elif resolved_document_type == DOCUMENT_TYPE_UNSTRUCTURED:
        if show_steps:
            logger.info("[step] unstructured section generation")
        sections, outline_meta = build_unstructured_sections(
            page_text_by_page=page_text_by_page,
            model=selected_model,
            use_llm=use_llm_effective,
            max_token_num_each_node=max_token_num_each_node,
            max_page_num_each_node=max_page_num_each_node,
        )
        structure = build_structure_from_sections(sections, include_summary_fields=summary_enabled)
        segment_meta = {
            "heading_anchor_match_rate": 1.0 if structure else 0.0,
            "overlap_violations": 0,
            "duplicated_block_rate": 0.0,
        }
    else:  # transcript
        if show_steps:
            logger.info("[step] transcript topic extraction")
        sections, outline_meta = build_transcript_sections(
            page_text_by_page=page_text_by_page,
            model=str(selected_model),
            max_token_num_each_node=max_token_num_each_node,
            max_page_num_each_node=max_page_num_each_node,
            prompt_template=transcript_topic_prompt_template,
            include_topic_summary=summary_enabled,
        )
        structure = build_structure_from_sections(sections, include_summary_fields=summary_enabled)
        segment_meta = {
            "heading_anchor_match_rate": 1.0 if structure else 0.0,
            "overlap_violations": 0,
            "duplicated_block_rate": 0.0,
        }

    transcript_payload = None
    if resolved_document_type == DOCUMENT_TYPE_TRANSCRIPT and isinstance(outline_meta.get("transcript"), dict):
        transcript_payload = dict(outline_meta.get("transcript", {}))
        outline_meta = {k: v for k, v in outline_meta.items() if k != "transcript"}

    limit_meta = enforce_node_limits(
        structure,
        max_token_num_each_node=max(1, int(max_token_num_each_node)),
        max_page_num_each_node=max(1, int(max_page_num_each_node)),
        include_summary_fields=summary_enabled,
    )
    segment_meta = {**segment_meta, **limit_meta}
    if show_details:
        logger.debug("[segmentation] metrics=%s", segment_meta)
        logger.debug("[segmentation] top_level_titles=%s", _summarize_top_level_titles(structure))

    if show_steps:
        logger.info("[step] summaries")
    _run_async(
        _summarize(
            structure=structure,
            model=selected_model,
            use_llm=use_llm_effective,
            logger=logger,
            include_summary_fields=summary_enabled,
        )
    )

    if not summary_enabled:
        _drop_summary_fields(structure)

    flat_nodes = _flatten(structure)
    summary_coverage = (
        sum(1 for n in flat_nodes if (n.get("summary") or "").strip()) / max(1, len(flat_nodes))
        if summary_enabled
        else 0.0
    )
    logger.info("[result] nodes=%d summary_coverage=%.2f", len(flat_nodes), summary_coverage)
    if show_details:
        logger.debug("[cost] usage=%s", get_usage_summary())

    navexa_document = {
        "doc_id": _slug(pdf_path),
        "doc_name": Path(pdf_path).stem,
        "pages": {"count": total_pages},
        "pipeline_version": PIPELINE_VERSION,
        "source": {
            "pdf_path": os.path.abspath(pdf_path),
            "sha256": _sha256(pdf_path),
            "total_pages": total_pages,
        },
        "cost": get_usage_summary(),
        "pipeline": {
            **outline_meta,
            "document_type": resolved_document_type,
            "summary_enabled": summary_enabled,
            "max_token_num_each_node": int(max_token_num_each_node),
            "max_page_num_each_node": int(max_page_num_each_node),
            "llm_required": llm_required,
            "requested_mode": resolved_mode,
            "effective_mode": MODE_LLM if use_llm_effective else MODE_NO_LLM,
            "verbosity": resolved_verbosity,
            "parser_model": resolved_parser_model,
            "output_format": resolved_output_format,
        },
        "structure": structure,
    }
    if transcript_payload is not None:
        navexa_document["transcript"] = transcript_payload

    validation = {
        "generated_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "pipeline_version": PIPELINE_VERSION,
        "metrics": {
            **validation_metrics(structure),
            **segment_meta,
            "summary_coverage": summary_coverage,
        },
    }

    compat = navexa_to_legacy_compat(navexa_document)
    logger.info("[done] indexing complete")
    return navexa_document, validation, compat


def write_outputs(
    out_dir: str,
    navexa_document: Dict[str, Any],
    validation_report: Dict[str, Any],
    compat_document: Dict[str, Any] | None = None,
    with_validation: bool = False,
    with_compat: bool = False,
) -> Dict[str, str]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    tree_path = out / "tree_navexa.json"
    tree_path.write_text(json.dumps(navexa_document, indent=2, ensure_ascii=False), encoding="utf-8")
    result = {"tree_navexa": str(tree_path)}

    if with_validation:
        validation_path = out / "validation_report.json"
        validation_path.write_text(json.dumps(validation_report, indent=2, ensure_ascii=False), encoding="utf-8")
        result["validation_report"] = str(validation_path)

    if with_compat and compat_document is not None:
        compat_path = out / "tree_legacy_compat.json"
        compat_path.write_text(json.dumps(compat_document, indent=2, ensure_ascii=False), encoding="utf-8")
        result["tree_legacy_compat"] = str(compat_path)

    return result


def run_cli() -> None:
    parser = argparse.ArgumentParser(description="Run Navexa indexing")
    parser.add_argument("--pdf", required=True, help="Input PDF path")
    parser.add_argument("--out-dir", required=True, help="Output directory")
    parser.add_argument("--model", default=None, help="Model/deployment name (llm mode)")
    parser.add_argument("--mode", default=os.getenv("NAVEXA_MODE", MODE_NO_LLM), help="llm | no-llm")
    parser.add_argument(
        "--document-type",
        default=os.getenv("NAVEXA_DOCUMENT_TYPE", DOCUMENT_TYPE_STRUCTURED),
        choices=[DOCUMENT_TYPE_STRUCTURED, DOCUMENT_TYPE_SEMI_STRUCTURED, DOCUMENT_TYPE_UNSTRUCTURED, DOCUMENT_TYPE_TRANSCRIPT],
        help="structured | semi_structured | unstructured | transcript",
    )
    parser.add_argument(
        "--parser-model",
        default=os.getenv("NAVEXA_PARSER_MODEL", "docling"),
        help="Parser backend model (currently: docling)",
    )
    parser.add_argument(
        "--output-format",
        default=os.getenv("NAVEXA_DOCLING_OUTPUT_FORMAT", "markdown"),
        choices=["markdown", "text"],
        help="Docling parsed output format used for heading/tree parsing",
    )
    parser.add_argument(
        "--verbose",
        default=os.getenv("NAVEXA_VERBOSE", VERBOSITY_MEDIUM),
        help="1/2/3 or low/medium/high",
    )
    parser.add_argument(
        "--max-token-num-each-node",
        type=int,
        default=int(os.getenv("NAVEXA_MAX_TOKEN_NUM_EACH_NODE", "12000")),
        help="Maximum tokens allowed per content node before split",
    )
    parser.add_argument(
        "--max-page-num-each-node",
        type=int,
        default=int(os.getenv("NAVEXA_MAX_PAGE_NUM_EACH_NODE", "8")),
        help="Maximum pages allowed per content node before split",
    )
    parser.add_argument(
        "--if-add-node-summary",
        default=os.getenv("NAVEXA_IF_ADD_NODE_SUMMARY", "yes"),
        choices=["yes", "no"],
        help="yes to include node summaries; no for text-only nodes",
    )
    parser.add_argument("--with-validation", action="store_true", help="Also write validation_report.json")
    parser.add_argument("--with-compat", action="store_true", help="Also write tree_legacy_compat.json")
    args = parser.parse_args()

    navexa_document, validation, compat = build_navexa_document(
        args.pdf,
        model=args.model,
        mode=args.mode,
        verbosity=args.verbose,
        parser_model=args.parser_model,
        output_format=args.output_format,
        document_type=args.document_type,
        max_token_num_each_node=args.max_token_num_each_node,
        max_page_num_each_node=args.max_page_num_each_node,
        if_add_node_summary=args.if_add_node_summary,
    )
    outputs = write_outputs(
        out_dir=args.out_dir,
        navexa_document=navexa_document,
        validation_report=validation,
        compat_document=compat,
        with_validation=args.with_validation,
        with_compat=args.with_compat,
    )

    if len(outputs) == 1:
        print(outputs["tree_navexa"])
    else:
        print(json.dumps(outputs, indent=2))
