from __future__ import annotations

from typing import Any

import gradio as gr

from ocr_studio.config import (
    EXAMPLES_DIR,
    LOCALE_EN,
    LOCALE_FA,
    MAX_PDF_PAGES,
    MAX_UPLOAD_MB,
    OCR_API_KEY,
)
from ocr_studio.i18n import language_choices, locale_choices, mode_choices, normalize_locale, t
from ocr_studio.pipeline import history_headers

CSS = """
@import url("https://fonts.googleapis.com/css2?family=IBM+Plex+Sans:wght@400;500;600&family=Vazirmatn:wght@400;600;700&display=swap");

:root {
  --ocr-bg: #f4f1ea;
  --ocr-ink: #1c1915;
  --ocr-muted: #5c564c;
  --ocr-line: #d8d0c3;
  --ocr-card: #fffcf7;
  --ocr-accent: #1f4f4a;
}

.gradio-container {
  font-family: "IBM Plex Sans", "Vazirmatn", sans-serif !important;
  max-width: 1180px !important;
  margin: 0 auto !important;
}

#ocr-shell {
  background: var(--ocr-card);
  border: 1px solid var(--ocr-line);
  border-radius: 22px;
  padding: 28px 28px 18px;
}

.rtl-ui {
  direction: rtl;
  font-family: "Vazirmatn", "IBM Plex Sans", sans-serif;
}

.ltr-ui {
  direction: ltr;
}

#hero h1 {
  font-size: 2rem;
  letter-spacing: -0.03em;
  margin-bottom: 0.35rem;
}

#hero p {
  color: var(--ocr-muted);
  font-size: 1.02rem;
  line-height: 1.7;
}

.note-grid {
  display: grid;
  grid-template-columns: 1fr;
  gap: 10px;
  margin: 16px 0 8px;
}

.note {
  border: 1px solid var(--ocr-line);
  background: #faf7f1;
  border-radius: 14px;
  padding: 12px 14px;
  color: var(--ocr-muted);
  font-size: 0.92rem;
  line-height: 1.6;
}

#run-btn, #rebuild-btn {
  min-height: 48px;
}

#result-text textarea {
  font-family: "Vazirmatn", "IBM Plex Sans", sans-serif;
  font-size: 1.02rem;
  line-height: 1.85;
}

footer, .built-with {
  display: none !important;
}

@media (max-width: 720px) {
  #ocr-shell { padding: 18px 14px 10px; }
  #hero h1 { font-size: 1.55rem; }
}
"""

THEME = gr.themes.Soft(
    font=[gr.themes.GoogleFont("IBM Plex Sans"), "Vazirmatn", "sans-serif"],
    primary_hue="teal",
    neutral_hue="stone",
)


def _notes(locale: str) -> str:
    return (
        '<div class="note-grid">'
        f'<div class="note">{t(locale, "note.limits", max_mb=MAX_UPLOAD_MB, max_pages=MAX_PDF_PAGES)}</div>'
        f'<div class="note">{t(locale, "note.privacy")}</div>'
        f'<div class="note">{t(locale, "note.quality")}</div>'
        "</div>"
    )


def _example_rows() -> list[list[object]]:
    rows: list[list[object]] = []
    mapping = [
        ("english-print.png", "en", "document"),
        ("persian-print.png", "fa", "document"),
        ("mixed-print.png", "fa+en", "precise"),
    ]
    for name, language, mode in mapping:
        path = EXAMPLES_DIR / name
        if path.exists():
            rows.append([str(path), language, mode])
    return rows


def _warning_markdown(warnings: list[str]) -> str:
    if not warnings:
        return ""
    return "\n".join(f"- {item}" for item in warnings)


def locale_updates(
    locale: str,
    language_value: str | None = None,
    mode_value: str | None = None,
) -> list[Any]:
    locale = normalize_locale(locale)
    require_key = bool(OCR_API_KEY)
    return [
        gr.update(value=t(locale, "intro")),
        gr.update(value=_notes(locale)),
        gr.update(label=t(locale, "image")),
        gr.update(label=t(locale, "files")),
        gr.update(
            label=t(locale, "language"),
            choices=language_choices(locale),
            value=language_value or "auto",
        ),
        gr.update(
            label=t(locale, "mode"),
            choices=mode_choices(locale),
            value=mode_value or "document",
        ),
        gr.update(label=t(locale, "deskew")),
        gr.update(label=t(locale, "accurate")),
        gr.update(label=t(locale, "api_key"), visible=require_key),
        gr.update(value=t(locale, "run")),
        gr.update(value=t(locale, "rebuild")),
        gr.update(label=t(locale, "text")),
        gr.update(label=t(locale, "preview")),
        gr.update(label=t(locale, "pdf")),
        gr.update(label=t(locale, "txt")),
        gr.update(label=t(locale, "md")),
        gr.update(label=t(locale, "json")),
        gr.update(label=t(locale, "docx")),
        gr.update(label=t(locale, "zip")),
        gr.update(label=t(locale, "history"), headers=history_headers(locale)),
        gr.update(label=t(locale, "ui_language"), choices=locale_choices(locale), value=locale),
        gr.update(value=t(locale, "footer")),
        gr.update(elem_classes=["rtl-ui"] if locale == LOCALE_FA else ["ltr-ui"]),
        locale,
    ]


def build_demo(run_fn, rebuild_fn, health_fn) -> gr.Blocks:
    locale0 = LOCALE_EN
    require_key = bool(OCR_API_KEY)

    with gr.Blocks(title="Document OCR") as demo:
        job_state = gr.State({})
        history_state = gr.State([])
        locale_state = gr.State(locale0)
        saved_locale = gr.BrowserState(locale0)

        with gr.Column(elem_id="ocr-shell", elem_classes=["ltr-ui"]) as shell:
            intro = gr.Markdown(t(locale0, "intro"), elem_id="hero")
            notes = gr.HTML(_notes(locale0))
            with gr.Row():
                locale_dd = gr.Dropdown(
                    label=t(locale0, "ui_language"),
                    choices=locale_choices(locale0),
                    value=locale0,
                    scale=1,
                )
            with gr.Row():
                with gr.Column(scale=5):
                    image = gr.Image(
                        label=t(locale0, "image"),
                        type="pil",
                        sources=["upload", "webcam", "clipboard"],
                        height=320,
                    )
                    files = gr.File(
                        label=t(locale0, "files"),
                        file_count="multiple",
                        file_types=["image", ".pdf", ".heic", ".heif"],
                        type="filepath",
                    )
                    language = gr.Dropdown(
                        label=t(locale0, "language"),
                        choices=language_choices(locale0),
                        value="auto",
                    )
                    mode = gr.Dropdown(
                        label=t(locale0, "mode"),
                        choices=mode_choices(locale0),
                        value="document",
                    )
                    deskew = gr.Checkbox(label=t(locale0, "deskew"), value=True)
                    high_accuracy = gr.Checkbox(label=t(locale0, "accurate"), value=False)
                    api_key = gr.Textbox(
                        label=t(locale0, "api_key"),
                        type="password",
                        visible=require_key,
                    )
                    run = gr.Button(t(locale0, "run"), variant="primary", elem_id="run-btn")
                    rebuild = gr.Button(t(locale0, "rebuild"), elem_id="rebuild-btn")
                with gr.Column(scale=7):
                    status = gr.Markdown(t(locale0, "status.wait"))
                    warnings = gr.Markdown("")
                    preview = gr.Gallery(
                        label=t(locale0, "preview"),
                        columns=2,
                        height=280,
                        preview=True,
                    )
                    text = gr.Textbox(
                        label=t(locale0, "text"),
                        lines=12,
                        elem_id="result-text",
                        rtl=False,
                        buttons=["copy"],
                    )
                    with gr.Row():
                        pdf = gr.File(label=t(locale0, "pdf"))
                        txt = gr.File(label=t(locale0, "txt"))
                    with gr.Row():
                        markdown = gr.File(label=t(locale0, "md"))
                        json_file = gr.File(label=t(locale0, "json"))
                    with gr.Row():
                        docx = gr.File(label=t(locale0, "docx"))
                        archive = gr.File(label=t(locale0, "zip"))
                    history = gr.Dataframe(
                        label=t(locale0, "history"),
                        headers=history_headers(locale0),
                        wrap=True,
                        interactive=False,
                    )

            examples = _example_rows()
            if examples:
                gr.Examples(
                    label="Examples",
                    examples=examples,
                    inputs=[image, language, mode],
                    cache_examples=False,
                )
            footer = gr.Markdown(t(locale0, "footer"))

        locale_outputs = [
            intro,
            notes,
            image,
            files,
            language,
            mode,
            deskew,
            high_accuracy,
            api_key,
            run,
            rebuild,
            text,
            preview,
            pdf,
            txt,
            markdown,
            json_file,
            docx,
            archive,
            history,
            locale_dd,
            footer,
            shell,
            locale_state,
        ]

        def switch_locale(locale, saved, language_value, mode_value):
            chosen = normalize_locale(locale or saved or locale0)
            return locale_updates(chosen, language_value, mode_value) + [chosen]

        locale_dd.change(
            fn=switch_locale,
            inputs=[locale_dd, saved_locale, language, mode],
            outputs=locale_outputs + [saved_locale],
            api_visibility="private",
        )
        demo.load(
            fn=lambda saved: locale_updates(normalize_locale(saved)),
            inputs=[saved_locale],
            outputs=locale_outputs,
            api_visibility="private",
        )

        run_outputs = [
            status,
            warnings,
            text,
            preview,
            pdf,
            txt,
            markdown,
            json_file,
            docx,
            archive,
            job_state,
            history_state,
            history,
        ]

        run.click(
            fn=run_fn,
            inputs=[
                image,
                files,
                language,
                mode,
                deskew,
                high_accuracy,
                api_key,
                locale_state,
                history_state,
            ],
            outputs=run_outputs,
            api_name="ocr",
            concurrency_limit=1,
        )
        rebuild.click(
            fn=rebuild_fn,
            inputs=[text, job_state, locale_state, history_state],
            outputs=run_outputs,
            api_name="rebuild",
            concurrency_limit=1,
        )

        health_in = gr.Button(visible=False)
        health_out = gr.Textbox(visible=False)
        health_in.click(fn=health_fn, outputs=health_out, api_name="health")

    return demo