import asyncio
import hashlib
import os
import re
import time
from collections import OrderedDict
from typing import Any, Optional

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

APP_DIR = os.path.dirname(os.path.abspath(__file__))
STATIC_DIR = os.path.join(APP_DIR, "static")

API_BASE = os.getenv("TRANSLATE_API_BASE", "https://dichmay.itrithuc.vn:1721").rstrip("/")
VERIFY_SSL = os.getenv("TRANSLATE_VERIFY_SSL", "true").lower() not in {"0", "false", "no"}
MAX_CHARS = int(os.getenv("MAX_CHARS", "12000"))
CACHE_TTL_SEC = int(os.getenv("CACHE_TTL_SEC", str(24 * 3600)))
CACHE_MAX_ITEMS = int(os.getenv("CACHE_MAX_ITEMS", "4096"))
TRANSLATE_TIMEOUT_SEC = float(os.getenv("TRANSLATE_TIMEOUT_SEC", "65"))
MAX_UPSTREAM_CONCURRENCY = int(os.getenv("MAX_UPSTREAM_CONCURRENCY", "2"))
POLL_MIN_DELAY_SEC = float(os.getenv("POLL_MIN_DELAY_SEC", "0.35"))
POLL_MAX_DELAY_SEC = float(os.getenv("POLL_MAX_DELAY_SEC", "4.0"))

QUEUE_WAIT_RE = re.compile(r"(?:ước\s*tính|uoc\s*tinh|estimated|wait|chờ|cho)[^0-9]{0,24}(\d{1,4})\s*(?:s|sec|second|giây|giay)?", re.IGNORECASE)

LANGS = {"vi", "en", "zh", "lo", "km"}

class TTLCache:
    def __init__(self, max_items: int = 4096, ttl: int = 86400) -> None:
        self.max_items = max_items
        self.ttl = ttl
        self._data: OrderedDict[str, tuple[float, str]] = OrderedDict()
        self._lock = asyncio.Lock()

    async def get(self, key: str) -> Optional[str]:
        now = time.time()
        async with self._lock:
            hit = self._data.get(key)
            if not hit:
                return None
            ts, value = hit
            if now - ts > self.ttl:
                self._data.pop(key, None)
                return None
            self._data.move_to_end(key)
            return value

    async def set(self, key: str, value: str) -> None:
        async with self._lock:
            self._data[key] = (time.time(), value)
            self._data.move_to_end(key)
            while len(self._data) > self.max_items:
                self._data.popitem(last=False)

    async def size(self) -> int:
        async with self._lock:
            return len(self._data)

cache = TTLCache(max_items=CACHE_MAX_ITEMS, ttl=CACHE_TTL_SEC)
inflight: dict[str, asyncio.Future[str]] = {}
inflight_lock = asyncio.Lock()
upstream_semaphore = asyncio.Semaphore(MAX_UPSTREAM_CONCURRENCY)

app = FastAPI(title="Bila Translator Realtime Space", version="4.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

class TranslateRequest(BaseModel):
    sourceText: str = Field(..., min_length=0, max_length=MAX_CHARS)
    sourceLang: str = Field(..., min_length=2, max_length=8)
    targetLang: str = Field(..., min_length=2, max_length=8)
    forceRefresh: bool = False

class TranslateResponse(BaseModel):
    ok: bool
    sourceLang: str
    targetLang: str
    text: str
    cached: bool = False
    ms: int
    error: Optional[str] = None

@app.get("/")
async def index() -> FileResponse:
    return FileResponse(os.path.join(STATIC_DIR, "index.html"))

@app.get("/health")
async def health() -> dict[str, Any]:
    return {
        "ok": True,
        "apiBase": API_BASE,
        "cacheItems": await cache.size(),
        "maxChars": MAX_CHARS,
        "maxUpstreamConcurrency": MAX_UPSTREAM_CONCURRENCY,
        "translateTimeoutSec": TRANSLATE_TIMEOUT_SEC,
    }

@app.get("/api/config")
async def config() -> dict[str, Any]:
    return {
        "langs": ["vi", "cvnss", "en", "zh", "lo", "km"],
        "maxChars": MAX_CHARS,
        "cacheTtlSec": CACHE_TTL_SEC,
    }

def normalize_lang(lang: str) -> str:
    lang = (lang or "").strip().lower()
    aliases = {"cn": "zh", "zh-cn": "zh", "zh_hans": "zh", "lao": "lo", "kh": "km", "vi-vn": "vi"}
    return aliases.get(lang, lang)

def cache_key(src_lang: str, target_lang: str, text: str) -> str:
    digest = hashlib.sha256(text.encode("utf-8")).hexdigest()
    return f"v4|{src_lang}|{target_lang}|{digest}"

async def fetch_json(client: httpx.AsyncClient, method: str, url: str, **kwargs: Any) -> Any:
    response = await client.request(method, url, **kwargs)
    response.raise_for_status()
    if not response.content:
        return {}
    return response.json()

def _deep_values(obj: Any) -> list[Any]:
    values: list[Any] = []
    if isinstance(obj, dict):
        for value in obj.values():
            values.append(value)
            values.extend(_deep_values(value))
    elif isinstance(obj, list):
        for value in obj:
            values.append(value)
            values.extend(_deep_values(value))
    return values

def extract_queue_wait_seconds(payload: Any) -> Optional[float]:
    """Read common upstream queue hints without depending on one exact schema."""
    if not isinstance(payload, dict):
        return None
    direct_names = {
        "estimatedWait", "estimated_wait", "estimatedWaitSec", "estimated_wait_sec",
        "estimatedWaitingSeconds", "waitingSeconds", "waitSeconds", "queueWait",
        "queue_wait", "queueWaitSec", "queue_wait_sec", "eta", "etaSec", "eta_sec",
    }
    for name in direct_names:
        value = payload.get(name)
        if isinstance(value, (int, float)) and value > 0:
            return float(value)
        if isinstance(value, str) and value.strip().isdigit():
            return float(value.strip())
    for value in _deep_values(payload):
        if isinstance(value, str):
            m = QUEUE_WAIT_RE.search(value)
            if m:
                return float(m.group(1))
    return None

async def poll_until_done(client: httpx.AsyncClient, history_id: str, task_id: str) -> str:
    # Queue-aware polling: do not hammer the upstream queue.  When the upstream
    # reports an estimated wait, sleep close to that hint before polling again.
    started = time.monotonic()
    delay = POLL_MIN_DELAY_SEC

    while time.monotonic() - started <= TRANSLATE_TIMEOUT_SEC:
        await asyncio.sleep(delay)
        url = f"{API_BASE}/translation-history/get-single"
        data = await fetch_json(
            client,
            "GET",
            url,
            params={"translationHistoryId": history_id, "taskId": task_id, "_t": int(time.time() * 1000)},
        )
        payload = data.get("data", data) if isinstance(data, dict) else {}
        status = str(payload.get("status", "")).lower()
        result_url = payload.get("resultUrl") or payload.get("result_url")
        if status == "translated" and result_url:
            return str(result_url)
        if status in {"failed", "error"}:
            raise RuntimeError("Upstream translation failed")

        queue_wait = extract_queue_wait_seconds(payload)
        if queue_wait:
            delay = max(POLL_MIN_DELAY_SEC, min(float(queue_wait) * 0.85, POLL_MAX_DELAY_SEC))
        else:
            delay = min(delay * 1.55, POLL_MAX_DELAY_SEC)

    raise TimeoutError("Translation timeout")

def extract_text(payload: Any) -> str:
    if not isinstance(payload, dict):
        return ""
    candidates = [
        payload.get("target_text"),
        payload.get("targetText"),
        payload.get("text"),
    ]
    data = payload.get("data")
    if isinstance(data, dict):
        candidates.extend([data.get("target_text"), data.get("targetText"), data.get("text")])
    for item in candidates:
        if isinstance(item, str) and item.strip():
            return item.strip()
    return ""

async def translate_upstream(text: str, source_lang: str, target_lang: str) -> str:
    # Protect the upstream translation queue.  Browser aborts do not necessarily
    # cancel a task that was already created upstream, so the Space must limit
    # how many fresh upstream jobs it creates at the same time.
    async with upstream_semaphore:
        timeout = httpx.Timeout(connect=8, read=max(TRANSLATE_TIMEOUT_SEC + 10, 75), write=8, pool=8)
        async with httpx.AsyncClient(timeout=timeout, verify=VERIFY_SSL, follow_redirects=True) as client:
            init = await fetch_json(
                client,
                "POST",
                f"{API_BASE}/translate",
                json={"sourceText": text, "sourceLang": source_lang, "targetLang": target_lang},
                headers={"Content-Type": "application/json"},
            )
            payload = init.get("data", init) if isinstance(init, dict) else {}
            task_id = payload.get("taskId")
            history_id = payload.get("translationHistoryId") or payload.get("translationHitoryId")
            if not task_id or not history_id:
                raise RuntimeError("Upstream did not return taskId/historyId")

            result_url = await poll_until_done(client, str(history_id), str(task_id))
            result_url = str(result_url).lstrip("/")
            final = await fetch_json(client, "GET", f"{API_BASE}/{result_url}", params={"_t": int(time.time() * 1000)})
            translated = extract_text(final)
            if not translated:
                raise RuntimeError("Upstream returned empty translation")
            return translated

async def translate_dedupe(text: str, source_lang: str, target_lang: str, key: str) -> str:
    async with inflight_lock:
        fut = inflight.get(key)
        if fut is None:
            loop = asyncio.get_running_loop()
            fut = loop.create_task(translate_upstream(text, source_lang, target_lang))
            inflight[key] = fut
    try:
        return await fut
    finally:
        async with inflight_lock:
            if inflight.get(key) is fut:
                inflight.pop(key, None)

@app.post("/api/translate", response_model=TranslateResponse)
async def translate(req: TranslateRequest) -> TranslateResponse:
    t0 = time.perf_counter()
    source_lang = normalize_lang(req.sourceLang)
    target_lang = normalize_lang(req.targetLang)
    text = req.sourceText.strip()

    if not text or source_lang == target_lang:
        return TranslateResponse(ok=True, sourceLang=source_lang, targetLang=target_lang, text=text, cached=True, ms=0)
    if source_lang == "cvnss" or target_lang == "cvnss":
        raise HTTPException(status_code=400, detail="CVNSS is converted in the browser for offline-speed sync")
    if source_lang not in LANGS or target_lang not in LANGS:
        raise HTTPException(status_code=400, detail=f"Unsupported language pair: {source_lang}->{target_lang}")

    key = cache_key(source_lang, target_lang, text)
    if not req.forceRefresh:
        hit = await cache.get(key)
        if hit is not None:
            return TranslateResponse(
                ok=True,
                sourceLang=source_lang,
                targetLang=target_lang,
                text=hit,
                cached=True,
                ms=int((time.perf_counter() - t0) * 1000),
            )

    try:
        translated = await translate_dedupe(text, source_lang, target_lang, key)
        await cache.set(key, translated)
        return TranslateResponse(
            ok=True,
            sourceLang=source_lang,
            targetLang=target_lang,
            text=translated,
            cached=False,
            ms=int((time.perf_counter() - t0) * 1000),
        )
    except Exception as exc:
        return JSONResponse(
            status_code=502,
            content={
                "ok": False,
                "sourceLang": source_lang,
                "targetLang": target_lang,
                "text": "",
                "cached": False,
                "ms": int((time.perf_counter() - t0) * 1000),
                "error": str(exc),
            },
        )