from __future__ import annotations

import json
import os
import time
import traceback
from pathlib import Path

os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


def _load_dotenv() -> None:
    env_path = Path(__file__).resolve().parent / ".env"
    if not env_path.exists():
        return
    for raw in env_path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


_load_dotenv()
_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
if _token:
    os.environ["HF_TOKEN"] = _token
    os.environ["HUGGING_FACE_HUB_TOKEN"] = _token

import spaces

from ocr_studio.assets import ensure_assets
from ocr_studio.config import (
    APP_AUTH_PASSWORD,
    APP_AUTH_USER,
    APP_PORT,
    MODEL_ID,
    QUEUE_SIZE,
    ZERO_GPU_DURATION,
    ZERO_GPU_DURATION_CAP,
)
from ocr_studio.errors import OcrError
from ocr_studio.i18n import t
from ocr_studio.metrics import metrics
from ocr_studio.pipeline import OcrPipeline, append_history, history_rows
from ocr_studio.ui import CSS, THEME, build_demo, _warning_markdown

ensure_assets()

pipeline = OcrPipeline()
pipeline.warmup()


def _count_inputs(image, files) -> int:
    count = 1 if image is not None else 0
    if files is None or files == "":
        return max(1, count)
    if isinstance(files, (list, tuple)):
        return max(1, count + len(files))
    return max(1, count + 1)


def estimate_run(image, files, language, mode, deskew, high_accuracy, api_key, locale, history) -> int:
    count = _count_inputs(image, files)
    factor = 1.0
    mode_key = str(mode or "document")
    if high_accuracy:
        factor *= 1.75
    lowered = mode_key.lower()
    if "compare" in lowered:
        factor *= 1.9
    elif "precise" in lowered:
        factor *= 1.12
    elif "table" in lowered:
        factor *= 1.2
    return min(ZERO_GPU_DURATION_CAP, max(ZERO_GPU_DURATION, int(16 + count * 20 * factor)))


def _to_ui(result, locale, history):
    import gradio as gr

    updated = append_history(history, result)
    rtl = result.rtl
    text_update = gr.update(
        value=result.text,
        rtl=rtl,
        text_align="right" if rtl else "left",
    )
    return (
        result.status,
        _warning_markdown(result.warnings),
        text_update,
        result.annotated,
        result.pdf_path,
        result.text_path,
        result.markdown_path,
        result.json_path,
        result.docx_path,
        result.zip_path,
        result.job,
        updated,
        history_rows(updated, locale),
    )


@spaces.GPU(duration=estimate_run)
def run_ocr(image, files, language, mode, deskew, high_accuracy, api_key, locale, history):
    import gradio as gr

    started = time.perf_counter()
    try:
        result = pipeline.run(
            image,
            language,
            mode,
            files=files,
            deskew=bool(deskew),
            high_accuracy=bool(high_accuracy),
            locale=locale or "en",
            api_key=api_key,
        )
    except OcrError as exc:
        metrics.record_error()
        raise gr.Error(exc.localized(locale)) from exc
    except Exception as exc:
        metrics.record_error()
        traceback.print_exc()
        raise gr.Error(t(locale, "err.failed")) from exc
    elapsed_ms = int((time.perf_counter() - started) * 1000)
    metrics.record_success(result.page_count, elapsed_ms)
    return _to_ui(result, locale or "en", history)


def rebuild_outputs(text, job, locale, history):
    import gradio as gr

    try:
        result = pipeline.rebuild_from_text(job or {}, text or "", locale or "en")
    except OcrError as exc:
        raise gr.Error(exc.localized(locale)) from exc
    except Exception as exc:
        traceback.print_exc()
        raise gr.Error(t(locale, "err.failed")) from exc
    return _to_ui(result, locale or "en", history)


def health_status() -> str:
    snapshot = metrics.snapshot().as_dict()
    payload = {
        "status": "ok",
        "model": MODEL_ID,
        "model_loaded": pipeline.engine.model is not None,
        "metrics": snapshot,
        "uptime_sec": int(time.time() - metrics.started_at),
    }
    return json.dumps(payload, ensure_ascii=False)


demo = build_demo(run_ocr, rebuild_outputs, health_status)

if __name__ == "__main__":
    auth = None
    if APP_AUTH_USER and APP_AUTH_PASSWORD:
        auth = (APP_AUTH_USER, APP_AUTH_PASSWORD)
    launch_kwargs = {
        "server_name": os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
        "server_port": APP_PORT,
        "ssr_mode": False,
        "theme": THEME,
        "css": CSS,
        "auth": auth,
    }
    demo.queue(default_concurrency_limit=1, max_size=QUEUE_SIZE).launch(**launch_kwargs)