from __future__ import annotations

import hmac
import zipfile
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from PIL import Image

from ocr_studio.config import (
    HISTORY_LIMIT,
    LANGUAGE_KEYS,
    MODE_KEYS,
    OCR_API_KEY,
    PAGE_BREAK,
)
from ocr_studio.engine import BatchInference, PageInference, PaddleOcrVlEngine
from ocr_studio.errors import OcrError
from ocr_studio.export_docx import write_docx_file
from ocr_studio.export_json import write_json_file
from ocr_studio.export_markdown import write_markdown_file
from ocr_studio.export_pdf import build_searchable_pdf
from ocr_studio.export_text import split_pages, write_text_file
from ocr_studio.i18n import LANGUAGE_LABEL_KEYS, MODE_LABEL_KEYS, STRINGS, t
from ocr_studio.image_io import collect_pages
from ocr_studio.jobs import new_job_dir
from ocr_studio.language import is_rtl_language, normalize_text, resolve_language
from ocr_studio.overlay import annotate_pages
from ocr_studio.preprocess import prepare_page
from ocr_studio.quality import assess_quality
from ocr_studio.spotting import TextSpan


@dataclass
class PreparedJob:
    pages: list[Image.Image]
    deskew_angles: list[float]
    truncated: bool
    source_page_count: int


@dataclass
class OcrOutputs:
    text: str
    language: str
    mode: str
    rtl: bool
    status: str
    warnings: list[str]
    annotated: list[Image.Image]
    pdf_path: str
    text_path: str
    markdown_path: str
    json_path: str
    docx_path: str
    zip_path: str
    page_count: int
    job: dict = field(default_factory=dict)


def _labels_for(key: str, table: dict[str, str]) -> set[str]:
    values = {key, table[key]}
    for bundle in STRINGS.values():
        values.add(bundle.get(table[key], ""))
    return {item for item in values if item}


def resolve_option(value: str, keys: tuple[str, ...], table: dict[str, str]) -> str:
    raw = (value or "").strip()
    if raw in keys:
        return raw
    for key in keys:
        if raw in _labels_for(key, table):
            return key
    raise OcrError("Invalid option selected.", key="err.option")


def resolve_language_option(value: str) -> str:
    return resolve_option(value, LANGUAGE_KEYS, LANGUAGE_LABEL_KEYS)


def resolve_mode_option(value: str) -> str:
    return resolve_option(value, MODE_KEYS, MODE_LABEL_KEYS)


def verify_api_key(provided: str | None) -> None:
    expected = OCR_API_KEY
    if not expected:
        return
    given = (provided or "").strip()
    if not given or not hmac.compare_digest(given, expected):
        raise OcrError("This server requires a valid API key.", key="err.api_key")


def _spans_payload(spans: list[TextSpan]) -> list[dict[str, object]]:
    payload: list[dict[str, object]] = []
    for span in spans:
        item: dict[str, object] = {"text": span.text}
        if span.box:
            item["box"] = [round(v, 2) for v in span.box]
        payload.append(item)
    return payload


def _save_page_images(pages: list[Image.Image], folder: Path) -> list[str]:
    paths: list[str] = []
    for index, page in enumerate(pages, start=1):
        path = folder / f"page-{index:02d}.jpg"
        page.convert("RGB").save(path, format="JPEG", quality=92)
        paths.append(str(path))
    return paths


def _zip_outputs(folder: Path, files: list[Path]) -> Path:
    archive = folder / "ocr-output.zip"
    with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as zipped:
        for path in files:
            if path.exists():
                zipped.write(path, arcname=path.name)
    return archive


def _combined_text(pages_text: list[str]) -> str:
    if len(pages_text) <= 1:
        return pages_text[0] if pages_text else ""
    return PAGE_BREAK.join(pages_text)


class OcrPipeline:
    def __init__(self, engine: PaddleOcrVlEngine | None = None) -> None:
        self.engine = engine or PaddleOcrVlEngine()

    def warmup(self) -> None:
        self.engine.load()

    def prepare_sources(
        self,
        image: Image.Image | str | Path | None,
        files: object,
        deskew: bool,
        high_accuracy: bool,
    ) -> PreparedJob:
        pages, source_count, truncated = collect_pages(image, files, accurate=high_accuracy)
        processed: list[Image.Image] = []
        angles: list[float] = []
        for page in pages:
            ready, angle = prepare_page(page, deskew=deskew)
            processed.append(ready)
            angles.append(angle)
        return PreparedJob(
            pages=processed,
            deskew_angles=angles,
            truncated=truncated,
            source_page_count=source_count,
        )

    def recognize(
        self,
        pages: list[Image.Image],
        language_key: str,
        mode_key: str,
        high_accuracy: bool,
    ) -> BatchInference:
        return self.engine.recognize_pages(pages, language_key, mode_key, high_accuracy)

    def export_result(
        self,
        prepared: PreparedJob,
        inference: BatchInference,
        language_key: str,
        mode_key: str,
        locale: str,
    ) -> OcrOutputs:
        page_texts: list[str] = []
        spans_by_page: list[list[TextSpan]] = []
        detected_language = language_key
        for page in inference.pages:
            language = resolve_language(language_key, page.display_text or page.alt_text)
            detected_language = language
            readable = page.alt_text if inference.compare and page.alt_text else page.display_text
            page_texts.append(normalize_text(readable, language))
            spans_by_page.append(page.spans)

        combined = _combined_text(page_texts)
        if not combined.strip():
            raise OcrError(
                "No readable text was found. Try a sharper, well-lit scan.",
                key="err.no_text",
            )

        warning_specs = assess_quality(
            combined,
            prepared.pages,
            spans_by_page,
            prepared.truncated,
            prepared.deskew_angles,
        )
        warnings = [t(locale, key, **params) for key, params in warning_specs]
        rtl = is_rtl_language(detected_language)
        annotated = annotate_pages(prepared.pages, spans_by_page)

        workdir = new_job_dir()
        page_paths = _save_page_images(prepared.pages, workdir)
        pdf_path = workdir / "ocr-scan.pdf"
        text_path = workdir / "ocr-text.txt"
        markdown_path = workdir / "ocr-document.md"
        json_path = workdir / "ocr-result.json"
        docx_path = workdir / "ocr-document.docx"

        build_searchable_pdf(prepared.pages, page_texts, pdf_path, spans=spans_by_page, rtl=rtl)
        write_text_file(combined, text_path)
        write_markdown_file(combined, markdown_path, detected_language, mode_key, len(prepared.pages))
        write_docx_file(combined, docx_path, detected_language)

        json_payload = {
            "engine": "PaddleOCR-VL-1.6",
            "language": detected_language,
            "mode": mode_key,
            "page_count": len(prepared.pages),
            "warnings": warnings,
            "text": combined,
            "pages": [
                {
                    "index": index,
                    "text": page_texts[index - 1],
                    "task": inference.pages[index - 1].task,
                    "tile_count": inference.pages[index - 1].tile_count,
                    "alt_text": inference.pages[index - 1].alt_text,
                    "spans": _spans_payload(spans_by_page[index - 1]),
                }
                for index in range(1, len(prepared.pages) + 1)
            ],
        }
        write_json_file(json_payload, json_path)
        zip_path = _zip_outputs(
            workdir,
            [pdf_path, text_path, markdown_path, json_path, docx_path],
        )

        char_count = len(combined)
        status = (
            f"{t(locale, 'run')} · {char_count} · {detected_language} · "
            f"{mode_key} · {len(prepared.pages)}p · PaddleOCR-VL-1.6"
        )
        if warnings:
            status = status + " · " + warnings[0]

        job_state = {
            "page_paths": page_paths,
            "spans": [
                [{"text": span.text, "box": span.box} for span in page_spans]
                for page_spans in spans_by_page
            ],
            "language": detected_language,
            "mode": mode_key,
            "rtl": rtl,
            "workdir": str(workdir),
        }
        return OcrOutputs(
            text=combined,
            language=detected_language,
            mode=mode_key,
            rtl=rtl,
            status=status,
            warnings=warnings,
            annotated=annotated,
            pdf_path=str(pdf_path),
            text_path=str(text_path),
            markdown_path=str(markdown_path),
            json_path=str(json_path),
            docx_path=str(docx_path),
            zip_path=str(zip_path),
            page_count=len(prepared.pages),
            job=job_state,
        )

    def rebuild_from_text(self, job: dict, text: str, locale: str) -> OcrOutputs:
        if not job or not job.get("page_paths"):
            raise OcrError(
                "Run extraction first, then edit the text and rebuild.",
                key="err.rebuild",
            )
        paths = [Path(item) for item in job["page_paths"]]
        missing = [path for path in paths if not path.exists()]
        if missing:
            raise OcrError(
                "Run extraction first, then edit the text and rebuild.",
                key="err.rebuild",
            )
        pages = [Image.open(path).convert("RGB") for path in paths]
        page_texts = split_pages(text or "", len(pages))
        language = job.get("language") or resolve_language("auto", text or "")
        mode_key = job.get("mode") or "document"
        rtl = bool(job.get("rtl")) or is_rtl_language(language)
        spans_by_page: list[list[TextSpan]] = []
        for page_spans in job.get("spans") or []:
            restored: list[TextSpan] = []
            for item in page_spans:
                box = item.get("box")
                restored.append(
                    TextSpan(
                        text=str(item.get("text") or ""),
                        box=tuple(box) if box else None,
                    )
                )
            spans_by_page.append(restored)
        while len(spans_by_page) < len(pages):
            spans_by_page.append([])

        prepared = PreparedJob(
            pages=pages,
            deskew_angles=[],
            truncated=False,
            source_page_count=len(pages),
        )
        inference = BatchInference(
            pages=[
                PageInference(
                    raw_text=page_texts[index],
                    display_text=page_texts[index],
                    spans=spans_by_page[index],
                    task="rebuild",
                )
                for index in range(len(pages))
            ]
        )
        return self.export_result(prepared, inference, language, mode_key, locale)

    def run(
        self,
        image: Image.Image | str | Path | None,
        language_label: str,
        mode_label: str,
        files: object = None,
        deskew: bool = True,
        high_accuracy: bool = False,
        locale: str = "en",
        api_key: str | None = None,
        recognize_fn=None,
    ) -> OcrOutputs:
        verify_api_key(api_key)
        language_key = resolve_language_option(language_label)
        mode_key = resolve_mode_option(mode_label)
        prepared = self.prepare_sources(image, files, deskew=deskew, high_accuracy=high_accuracy)
        worker = recognize_fn or self.recognize
        inference = worker(prepared.pages, language_key, mode_key, high_accuracy)
        return self.export_result(prepared, inference, language_key, mode_key, locale)


def append_history(history: list[dict] | None, result: OcrOutputs) -> list[dict]:
    items = list(history or [])
    items.append(
        {
            "time": datetime.now(timezone.utc).strftime("%H:%M:%S"),
            "pages": result.page_count,
            "language": result.language,
            "mode": result.mode,
            "chars": len(result.text),
            "preview": result.text.replace("\n", " ")[:180],
        }
    )
    return items[-HISTORY_LIMIT:]


def history_rows(history: list[dict] | None, locale: str) -> list[list[object]]:
    rows: list[list[object]] = []
    for item in history or []:
        rows.append(
            [
                item.get("time", ""),
                item.get("pages", 0),
                item.get("language", ""),
                item.get("mode", ""),
                item.get("chars", 0),
                item.get("preview", ""),
            ]
        )
    return rows


def history_headers(locale: str) -> list[str]:
    return [
        t(locale, "hist.time"),
        t(locale, "hist.pages"),
        t(locale, "hist.language"),
        t(locale, "hist.mode"),
        t(locale, "hist.chars"),
        t(locale, "hist.preview"),
    ]