from __future__ import annotations

import json
import warnings
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from ..pipeline.index_pdf import build_navexa_document
from .models import IndexResult, SaveResult


def _utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _default_output_dir(index_result: IndexResult) -> Path:
    source_pdf = (
        index_result.tree_navexa.get("source", {}).get("pdf_path")
        if isinstance(index_result.tree_navexa, dict)
        else None
    )
    if source_pdf:
        pdf_path = Path(source_pdf)
        return pdf_path.parent / f"{pdf_path.stem}_navexa_out"

    doc_id = str(index_result.tree_navexa.get("doc_id", "document")).strip() or "document"
    return Path.cwd() / f"{doc_id}_navexa_out"


def _build_index_result(
    *,
    pdf_path: str,
    model: Optional[str],
    mode: Optional[str],
    verbosity: Optional[str],
    parser_model: Optional[str],
    output_format: Optional[str],
    document_type: str,
    max_token_num_each_node: int,
    max_page_num_each_node: int,
    if_add_node_summary: str,
    semi_heading_prompt_template: Optional[str] = None,
    transcript_topic_prompt_template: Optional[str] = None,
) -> IndexResult:
    tree_navexa, validation_report, tree_legacy_compat = build_navexa_document(
        pdf_path=pdf_path,
        model=model,
        mode=mode,
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        document_type=document_type,
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
        semi_heading_prompt_template=semi_heading_prompt_template,
        transcript_topic_prompt_template=transcript_topic_prompt_template,
    )

    pipeline_meta = tree_navexa.get("pipeline", {}) if isinstance(tree_navexa, dict) else {}
    meta = {
        "mode": pipeline_meta.get("requested_mode"),
        "effective_mode": pipeline_meta.get("effective_mode"),
        "document_type": pipeline_meta.get("document_type"),
        "model": model or tree_navexa.get("cost", {}).get("pricing_model"),
        "cost": tree_navexa.get("cost", {}),
        "generated_at": _utc_now(),
    }

    return IndexResult(
        tree_navexa=tree_navexa,
        validation_report=validation_report,
        tree_legacy_compat=tree_legacy_compat,
        meta=meta,
    )


def index_structured_document_tree(
    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,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: str = "yes",
) -> IndexResult:
    return _build_index_result(
        pdf_path=pdf_path,
        model=model,
        mode=mode,
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        document_type="structured",
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
    )


def index_semi_structured_document_tree(
    pdf_path: str,
    *,
    model: Optional[str] = None,
    verbosity: Optional[str] = None,
    parser_model: Optional[str] = None,
    output_format: Optional[str] = None,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: str = "yes",
    semi_heading_prompt_template: Optional[str] = None,
) -> IndexResult:
    return _build_index_result(
        pdf_path=pdf_path,
        model=model,
        mode="llm",
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        document_type="semi_structured",
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
        semi_heading_prompt_template=semi_heading_prompt_template,
    )


def index_unstructured_document_tree(
    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,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: str = "yes",
) -> IndexResult:
    return _build_index_result(
        pdf_path=pdf_path,
        model=model,
        mode=mode,
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        document_type="unstructured",
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
    )


def index_transcript_document_tree(
    pdf_path: str,
    *,
    model: Optional[str] = None,
    verbosity: Optional[str] = None,
    parser_model: Optional[str] = None,
    output_format: Optional[str] = None,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: str = "yes",
    transcript_topic_prompt_template: Optional[str] = None,
) -> IndexResult:
    return _build_index_result(
        pdf_path=pdf_path,
        model=model,
        mode="llm",
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        document_type="transcript",
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
        transcript_topic_prompt_template=transcript_topic_prompt_template,
    )


def index_document_tree(
    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,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: str = "yes",
) -> IndexResult:
    warnings.warn(
        "index_document_tree is kept for backward compatibility. "
        "Use index_structured_document_tree/index_semi_structured_document_tree/"
        "index_unstructured_document_tree/index_transcript_document_tree instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    return index_structured_document_tree(
        pdf_path=pdf_path,
        model=model,
        mode=mode,
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
    )


def save_document_tree(
    index_result: IndexResult,
    out_dir: Optional[str] = None,
    save_mode: str = "explicit",
    *,
    write_tree: bool = True,
    write_validation: bool = False,
    write_compat: bool = False,
) -> SaveResult:
    del save_mode  # reserved for future behavior variants

    if not any([write_tree, write_validation, write_compat]):
        raise ValueError("At least one artifact must be enabled for writing.")

    target_dir = Path(out_dir) if out_dir else _default_output_dir(index_result)
    target_dir.mkdir(parents=True, exist_ok=True)

    paths = {}
    if write_tree:
        tree_path = target_dir / "tree_navexa.json"
        tree_path.write_text(json.dumps(index_result.tree_navexa, indent=2, ensure_ascii=False), encoding="utf-8")
        paths["tree_navexa"] = str(tree_path)

    if write_validation:
        validation_path = target_dir / "validation_report.json"
        validation_path.write_text(
            json.dumps(index_result.validation_report, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )
        paths["validation_report"] = str(validation_path)

    if write_compat:
        compat_path = target_dir / "tree_legacy_compat.json"
        compat_path.write_text(
            json.dumps(index_result.tree_legacy_compat, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )
        paths["tree_legacy_compat"] = str(compat_path)

    return SaveResult(out_dir=str(target_dir), paths=paths)


def index_and_save_document_tree(
    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,
    max_token_num_each_node: int = 12000,
    max_page_num_each_node: int = 8,
    if_add_node_summary: str = "yes",
    out_dir: Optional[str] = None,
    save_mode: str = "explicit",
    write_tree: bool = True,
    write_validation: bool = False,
    write_compat: bool = False,
) -> tuple[IndexResult, SaveResult]:
    index_result = index_structured_document_tree(
        pdf_path=pdf_path,
        model=model,
        mode=mode,
        verbosity=verbosity,
        parser_model=parser_model,
        output_format=output_format,
        max_token_num_each_node=max_token_num_each_node,
        max_page_num_each_node=max_page_num_each_node,
        if_add_node_summary=if_add_node_summary,
    )
    save_result = save_document_tree(
        index_result=index_result,
        out_dir=out_dir,
        save_mode=save_mode,
        write_tree=write_tree,
        write_validation=write_validation,
        write_compat=write_compat,
    )
    return index_result, save_result
