from __future__ import annotations

import json
import os
import inspect
import time
from concurrent.futures import TimeoutError
from pathlib import Path
from typing import Any

from gradio_client import Client, handle_file


class WorkerError(RuntimeError):
    pass


def normalize_endpoint(endpoint: str) -> str:
    value = (endpoint or "").strip().rstrip("/")
    if not value:
        return ""
    if not value.startswith(("http://", "https://")):
        value = "https://" + value
    return value


def _client(endpoint: str) -> Client:
    hf_token = os.getenv("HF_TOKEN", "").strip() or None
    if not hf_token:
        return Client(endpoint)

    signature = inspect.signature(Client.__init__)
    if "hf_token" in signature.parameters:
        return Client(endpoint, hf_token=hf_token)
    if "token" in signature.parameters:
        return Client(endpoint, token=hf_token)
    if "headers" in signature.parameters:
        return Client(endpoint, headers={"Authorization": f"Bearer {hf_token}"})
    return Client(endpoint)


def check_health(endpoint: str, timeout_seconds: float = 20.0) -> dict[str, Any]:
    endpoint = normalize_endpoint(endpoint)
    if not endpoint:
        raise WorkerError("Worker URL is not configured.")
    try:
        job = _client(endpoint).submit(api_name="/health")
        payload = job.result(timeout=timeout_seconds)
        if not isinstance(payload, dict):
            raise WorkerError("Worker health response is not a JSON object.")
        return payload
    except TimeoutError as exc:
        raise WorkerError(f"Worker health check timed out after {timeout_seconds:.0f}s.") from exc
    except Exception as exc:
        raise WorkerError(f"Worker health check failed: {exc}") from exc


def run_page(
    *,
    endpoint: str,
    model_id: str,
    page_path: Path,
    prompt: str,
    options: dict[str, Any],
    timeout_seconds: float,
) -> dict[str, Any]:
    endpoint = normalize_endpoint(endpoint)
    if not endpoint:
        raise WorkerError("Worker URL is not configured.")

    retry_delays = (0.0, 2.0, 6.0)
    last_error: Exception | None = None

    for attempt, delay in enumerate(retry_delays, start=1):
        if delay:
            time.sleep(delay)
        job = None
        try:
            job = _client(endpoint).submit(
                handle_file(str(page_path)),
                model_id,
                prompt or "",
                json.dumps(options, ensure_ascii=False),
                os.getenv("WORKER_API_TOKEN", "").strip(),
                api_name="/ocr",
            )
            payload = job.result(timeout=timeout_seconds)
            if not isinstance(payload, dict):
                raise WorkerError("Worker response is not a JSON object.")
            required = {"model", "text", "markdown", "metrics"}
            missing = sorted(required - payload.keys())
            if missing:
                raise WorkerError(f"Worker response is missing fields: {missing}")
            return payload
        except TimeoutError as exc:
            last_error = WorkerError(f"Worker request timed out after {timeout_seconds:.0f}s.")
            if job is not None:
                try:
                    job.cancel()
                except Exception:
                    pass
            if attempt >= len(retry_delays):
                break
        except (ValueError, OSError, WorkerError, Exception) as exc:
            last_error = exc
            if attempt >= len(retry_delays):
                break

    raise WorkerError(f"OCR request failed after retries: {last_error}")