"""
Document preview contract helpers (Option B + talkback UI metadata).
"""
from __future__ import annotations

import os
import re
import shutil
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
from urllib.parse import quote

PreviewMode = str  # pdf | image | text | html | none

OFFICE_PREVIEW_AS_PDF_EXT = frozenset({"doc", "docx"})
DEFAULT_PREVIEW_PDF_CACHE_DIR = "document_images_/preview_cache"
PAGE_PDF_CACHE_DIR = "document_images_/page_pdf_cache"


def document_media_type(filename: str) -> str:
    file_ext = os.path.splitext(filename or "")[1].lower()
    media_types = {
        ".pdf": "application/pdf",
        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ".doc": "application/msword",
        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        ".xls": "application/vnd.ms-excel",
        ".csv": "text/csv",
        ".txt": "text/plain",
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".bmp": "image/bmp",
        ".svg": "image/svg+xml",
        ".html": "text/html",
        ".htm": "text/html",
    }
    return media_types.get(file_ext, "application/octet-stream")


def normalize_file_type(file_type: str = "", filename: str = "") -> str:
    ft = (file_type or "").strip().lower().lstrip(".")
    if ft:
        return ft
    ext = os.path.splitext(filename or "")[1].lower().lstrip(".")
    return ext or "bin"


def resolve_document_preview_mode(file_type: str = "", filename: str = "") -> PreviewMode:
    ext = normalize_file_type(file_type, filename)
    if ext in OFFICE_PREVIEW_AS_PDF_EXT or ext == "pdf":
        return "pdf"
    if ext in ("png", "jpg", "jpeg", "gif", "webp", "bmp", "svg"):
        return "image"
    if ext in ("txt", "csv", "md", "json", "log"):
        return "text"
    if ext in ("html", "htm"):
        return "html"
    return "none"


def preview_content_type_for_contract(file_type: str = "", filename: str = "") -> str:
    """MIME type the preview endpoint will return (PDF for Office docs converted server-side)."""
    ext = normalize_file_type(file_type, filename)
    if ext in OFFICE_PREVIEW_AS_PDF_EXT:
        return "application/pdf"
    return document_media_type(filename or f"file.{ext}")


def _sanitize_html_for_story(html: str) -> str:
    """Remove pandoc CSS that PyMuPDF Story cannot parse."""
    html = re.sub(r"<style[^>]*>.*?</style>", "", html, flags=re.DOTALL | re.IGNORECASE)
    html = re.sub(
        r'<link[^>]*rel=["\']stylesheet["\'][^>]*>',
        "",
        html,
        flags=re.IGNORECASE,
    )
    return html


def convert_office_file_to_pdf(
    source_path: str,
    *,
    output_pdf_path: Optional[str] = None,
) -> Tuple[str, List[str]]:
    """
    Convert DOC/DOCX to PDF via pandoc HTML + PyMuPDF Story.
    Returns (pdf_path, temp_paths_to_cleanup). *output_pdf_path* is not listed for cleanup.
    """
    import pymupdf
    import pypandoc

    cleanup_paths: List[str] = []
    media_dir = tempfile.mkdtemp()
    cleanup_paths.append(media_dir)
    fd, html_path = tempfile.mkstemp(suffix=".html")
    os.close(fd)
    cleanup_paths.append(html_path)

    pypandoc.convert_file(
        source_path,
        "html",
        outputfile=html_path,
        extra_args=["--standalone", f"--extract-media={media_dir}"],
    )
    with open(html_path, "r", encoding="utf-8") as html_file:
        html_content = _sanitize_html_for_story(html_file.read())

    if output_pdf_path:
        pdf_path = output_pdf_path
        Path(pdf_path).parent.mkdir(parents=True, exist_ok=True)
    else:
        fd, pdf_path = tempfile.mkstemp(suffix=".pdf")
        os.close(fd)
        cleanup_paths.append(pdf_path)

    mediabox = pymupdf.paper_rect("a4")
    story = pymupdf.Story(html=html_content, archive=media_dir)
    writer = pymupdf.DocumentWriter(pdf_path)

    def rectfn(_rect_num: int, _filled) -> Tuple:
        return mediabox, mediabox, None

    story.write(writer, rectfn)
    writer.close()
    return pdf_path, cleanup_paths


def ensure_office_preview_pdf(
    source_path: str,
    file_id: str,
    cache_dir: str = DEFAULT_PREVIEW_PDF_CACHE_DIR,
) -> str:
    """
    Return a cached PDF path for Office preview, converting when missing or stale.
    """
    Path(cache_dir).mkdir(parents=True, exist_ok=True)
    cached = os.path.join(cache_dir, f"{file_id}.pdf")
    try:
        if os.path.isfile(cached) and os.path.getmtime(cached) >= os.path.getmtime(
            source_path
        ):
            return cached
    except OSError:
        pass
    convert_office_file_to_pdf(source_path, output_pdf_path=cached)
    return cached


def resolve_preview_file_response(
    source_path: str,
    original_filename: str,
    file_id: str,
    *,
    inline: bool,
    cache_dir: str = DEFAULT_PREVIEW_PDF_CACHE_DIR,
) -> Tuple[str, str, str]:
    """
    Resolve path, Content-Type, and Content-Disposition filename for preview/download.

    Inline Office files are served as converted PDF bytes.
    """
    name = (original_filename or os.path.basename(source_path)).strip()
    ext = normalize_file_type("", name)
    if inline and ext in OFFICE_PREVIEW_AS_PDF_EXT:
        pdf_path = ensure_office_preview_pdf(source_path, file_id, cache_dir=cache_dir)
        pdf_name = os.path.splitext(name)[0] + ".pdf"
        return pdf_path, "application/pdf", pdf_name
    return source_path, document_media_type(name), name


def build_document_preview_path(file_id: str, agent_id: str) -> str:
    fid = quote(str(file_id).strip(), safe="")
    aid = quote(str(agent_id).strip(), safe="")
    return f"/api/training/document/preview-file/{fid}?agent_id={aid}"


def build_document_page_download_path(
    file_id: str,
    agent_id: str,
    page_number: int,
) -> str:
    """Server-relative URL to download a single PDF page (attachment)."""
    fid = quote(str(file_id).strip(), safe="")
    aid = quote(str(agent_id).strip(), safe="")
    page = max(1, int(page_number))
    return (
        f"/api/training/document/download-page/{fid}"
        f"?agent_id={aid}&page_number={page}"
    )


def build_document_page_download_url(
    file_id: str,
    agent_id: str,
    page_number: int,
) -> str:
    from app.core.config import get_base_url

    base = get_base_url().rstrip("/")
    return f"{base}{build_document_page_download_path(file_id, agent_id, page_number)}"


def cached_page_pdf_path(
    file_id: str,
    page_number: int,
    cache_dir: str = PAGE_PDF_CACHE_DIR,
) -> str:
    page = max(1, int(page_number))
    return os.path.join(cache_dir, f"{file_id}_page_{page}.pdf")


def extract_and_cache_page_pdf(
    full_pdf_path: str,
    file_id: str,
    page_number: int,
    cache_dir: str = PAGE_PDF_CACHE_DIR,
) -> str:
    """Extract one page (1-based) from *full_pdf_path* into a cached single-page PDF."""
    import pymupdf

    page = max(1, int(page_number))
    out_path = cached_page_pdf_path(file_id, page, cache_dir=cache_dir)
    Path(cache_dir).mkdir(parents=True, exist_ok=True)

    src_mtime = os.path.getmtime(full_pdf_path) if os.path.isfile(full_pdf_path) else 0
    try:
        if os.path.isfile(out_path) and os.path.getmtime(out_path) >= src_mtime:
            return out_path
    except OSError:
        pass

    doc = pymupdf.open(full_pdf_path)
    try:
        if page > doc.page_count:
            page = doc.page_count
        if page < 1:
            raise ValueError("invalid page_number")
        new_doc = pymupdf.open()
        new_doc.insert_pdf(doc, from_page=page - 1, to_page=page - 1)
        new_doc.save(out_path)
        new_doc.close()
    finally:
        doc.close()
    return out_path


def build_document_page_image_path(
    file_id: str,
    agent_id: str,
    page_number: int,
) -> str:
    """Server-relative URL to render a single PDF page as PNG at request time."""
    fid = quote(str(file_id).strip(), safe="")
    aid = quote(str(agent_id).strip(), safe="")
    page = max(1, int(page_number))
    return (
        f"/api/training/document/page-image/{fid}"
        f"?agent_id={aid}&page_number={page}"
    )


def render_page_png_bytes(full_pdf_path: str, page_number: int) -> bytes:
    """Render one PDF page (1-based) to PNG bytes (no disk write)."""
    import pymupdf

    page = max(1, int(page_number))
    doc = pymupdf.open(full_pdf_path)
    try:
        if page > doc.page_count:
            page = doc.page_count
        if page < 1:
            raise ValueError("invalid page_number")
        pix = doc[page - 1].get_pixmap(matrix=pymupdf.Matrix(2, 2))
        return pix.tobytes("png")
    finally:
        doc.close()


def apply_page_image_url_to_mapping(
    target: Dict[str, Any],
    *,
    file_id: str,
    agent_id: str,
    page_number: int,
) -> None:
    """Set path / relevant_images to runtime page-image API URLs."""
    if not file_id or agent_id is None or not str(agent_id).strip():
        return
    pn = int(page_number)
    aid = str(agent_id).strip()
    rel = build_document_page_image_path(file_id, aid, pn)
    target["path"] = rel
    target["relevant_images"] = [
        {
            "path": rel,
            "page": pn,
            "type": "page_render",
            "is_matching": True,
        }
    ]


def build_page_images_from_api_paths(
    file_id: str,
    agent_id: str,
    pages: Sequence[int],
) -> Dict[str, List[Dict[str, Any]]]:
    """Build frontend ``page_images`` dict pointing at runtime page-image API."""
    out: Dict[str, List[Dict[str, Any]]] = {}
    fid = str(file_id).strip()
    aid = str(agent_id).strip()
    if not fid or not aid:
        return out
    for page in sorted({int(p) for p in pages if p is not None and int(p) >= 1}):
        rel = build_document_page_image_path(fid, aid, page)
        out[str(page)] = [
            {
                "path": rel,
                "type": "page_render",
                "is_matching": True,
            }
        ]
    return out


def apply_page_download_url_to_mapping(
    target: Dict[str, Any],
    *,
    file_id: str,
    agent_id: str,
    page_number: int,
    relative_url: bool = True,
) -> None:
    """Set page_download_url (and absolute variant) on a hit or find result row."""
    if not file_id or agent_id is None or not str(agent_id).strip():
        return
    pn = int(page_number)
    rel = build_document_page_download_path(file_id, str(agent_id).strip(), pn)
    target["page_number"] = pn
    target["page_download_url"] = rel
    if not relative_url:
        target["page_download_url_absolute"] = build_document_page_download_url(
            file_id, str(agent_id).strip(), pn
        )


def render_page_pdf_for_download(
    file_path: str,
    original_filename: str,
    file_id: str,
    page_number: int,
    *,
    cache_dir: str = DEFAULT_PREVIEW_PDF_CACHE_DIR,
    page_cache_dir: str = PAGE_PDF_CACHE_DIR,
) -> Optional[str]:
    """Resolve on-disk single-page PDF for email attachments and similar."""
    try:
        full_pdf_path, media_type, _ = resolve_preview_file_response(
            file_path,
            original_filename,
            file_id,
            inline=True,
            cache_dir=cache_dir,
        )
        if media_type != "application/pdf":
            return None
        return extract_and_cache_page_pdf(
            full_pdf_path,
            file_id,
            page_number,
            cache_dir=page_cache_dir,
        )
    except Exception:
        return None


def build_matching_page_urls(
    *,
    file_id: str,
    agent_id: str,
    pages: Sequence[int],
    relative_url: bool = True,
) -> List[Dict[str, Any]]:
    """List of {page, page_download_url} for @find matching_pages."""
    out: List[Dict[str, Any]] = []
    aid = str(agent_id).strip()
    fid = str(file_id).strip()
    for p in sorted({int(x) for x in pages if x is not None}):
        if p < 1:
            continue
        row: Dict[str, Any] = {
            "page": p,
            "page_download_url": build_document_page_download_path(fid, aid, p),
        }
        if not relative_url:
            row["page_download_url_absolute"] = build_document_page_download_url(fid, aid, p)
        out.append(row)
    return out


def build_document_preview_url(file_id: str, agent_id: str) -> str:
    """Absolute preview URL (for window.open without client base URL)."""
    from app.core.config import get_base_url

    base = get_base_url().rstrip("/")
    return f"{base}{build_document_preview_path(file_id, agent_id)}"


def build_document_preview_contract(
    *,
    file_id: str,
    agent_id: str,
    filename: str,
    file_type: str = "",
    relative_url: bool = True,
) -> Dict[str, Any]:
    """
    Talkback / search document preview metadata for the frontend.

    preview_url is server-relative by default (prefix with API base on the client).
    """
    ft = normalize_file_type(file_type, filename)
    name = (filename or "").strip() or f"document.{ft}"
    path = build_document_preview_path(file_id, agent_id)
    return {
        "file_id": str(file_id).strip(),
        "filename": name,
        "file_type": ft,
        "preview_url": path if relative_url else build_document_preview_url(file_id, agent_id),
        "preview_content_type": preview_content_type_for_contract(ft, name),
        "preview_mode": resolve_document_preview_mode(ft, name),
    }


def apply_preview_contract_to_mapping(
    target: Dict[str, Any],
    *,
    file_id: str,
    agent_id: str,
    filename: str = "",
    file_type: str = "",
    relative_url: bool = True,
) -> None:
    """Merge preview contract fields into an existing dict (chunk hit, find result, etc.)."""
    if not file_id or agent_id is None or not str(agent_id).strip():
        return
    contract = build_document_preview_contract(
        file_id=file_id,
        agent_id=str(agent_id).strip(),
        filename=filename or str(target.get("filename") or ""),
        file_type=file_type or str(target.get("file_type") or ""),
        relative_url=relative_url,
    )
    target.update(contract)
    if not relative_url:
        target["preview_url_absolute"] = contract.get("preview_url")


def collect_talkback_documents_from_file_results(
    file_results: Optional[List[Dict[str, Any]]],
) -> List[Dict[str, Any]]:
    """
    Build deduplicated ``details.documents`` list for talkback / ask responses
    from per-agent file_result payloads.
    """
    if not file_results:
        return []

    seen: set = set()
    out: List[Dict[str, Any]] = []

    def _add(
        file_id: Any,
        agent_id: Any,
        filename: str = "",
        file_type: str = "",
    ) -> None:
        if not file_id or agent_id is None or not str(agent_id).strip():
            return
        fid = str(file_id).strip()
        aid = str(agent_id).strip()
        key = (fid, aid)
        if key in seen:
            return
        seen.add(key)
        out.append(
            build_document_preview_contract(
                file_id=fid,
                agent_id=aid,
                filename=filename,
                file_type=file_type,
                relative_url=True,
            )
        )

    for fr in file_results:
        if not isinstance(fr, dict):
            continue
        agent_id = fr.get("agent_id")
        rtype = fr.get("document_result_type") or ""

        if rtype == "document_find":
            block = fr.get("document_find") or {}
            for row in block.get("results") or []:
                if not isinstance(row, dict):
                    continue
                _add(
                    row.get("file_id"),
                    row.get("agent_id") or agent_id,
                    str(row.get("filename") or ""),
                    str(row.get("file_type") or ""),
                )
            continue

        if rtype == "document_search":
            block = fr.get("document_search") or {}
            for hit in (block.get("chunk_hits") or []) + (block.get("results") or []):
                if not isinstance(hit, dict):
                    continue
                _add(
                    hit.get("file_id") or hit.get("id"),
                    hit.get("agent_id") or agent_id,
                    str(hit.get("filename") or ""),
                    str(hit.get("file_type") or ""),
                )
            continue

        if rtype == "document_analysis":
            block = fr.get("document_analysis") or fr
            fid = block.get("file_id") or fr.get("file_id")
            _add(
                fid,
                agent_id,
                str(block.get("filename") or ""),
                str(block.get("file_type") or ""),
            )

    return out


def collect_talkback_documents_from_details(
    details: Optional[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """
    Extract ``documents`` from API ``details`` (multi-agent file_results or ask-file blocks).
    """
    if not isinstance(details, dict):
        return []

    fr = details.get("file_results")
    if isinstance(fr, list) and fr:
        return collect_talkback_documents_from_file_results(fr)

    agent_id = details.get("agent_id")
    if agent_id is None:
        mids = details.get("matched_agent_ids")
        if isinstance(mids, list) and mids:
            agent_id = mids[0]

    rtype = details.get("document_result_type")
    if rtype == "document_find" and details.get("document_find"):
        return collect_talkback_documents_from_file_results(
            [
                {
                    "agent_id": agent_id,
                    "document_result_type": "document_find",
                    "document_find": details["document_find"],
                }
            ]
        )
    if rtype == "document_search" and details.get("document_search"):
        return collect_talkback_documents_from_file_results(
            [
                {
                    "agent_id": agent_id,
                    "document_result_type": "document_search",
                    "document_search": details["document_search"],
                }
            ]
        )
    if rtype == "document_analysis":
        block = details.get("document_analysis") or details
        fid = block.get("file_id") if isinstance(block, dict) else None
        if fid and agent_id is not None:
            return [
                build_document_preview_contract(
                    file_id=str(fid),
                    agent_id=str(agent_id),
                    filename=str(block.get("filename") or ""),
                    file_type=str(block.get("file_type") or ""),
                    relative_url=True,
                )
            ]
    return []
