import asyncio
from contextlib import suppress
import hashlib
import logging

from .config import ProviderConfig
from .models import Project
from .protection import protect
from .provider import Provider, ProviderError
from .storage import encode

log = logging.getLogger(__name__)
PROMPT_VERSION = 3


class Engine:
    def __init__(self, storage, provider_factory=Provider):
        self.storage = storage
        self.provider_factory = provider_factory
        self.wakeup = asyncio.Event()
        self.task = None

    def start(self):
        self.storage.recover()
        self.task = asyncio.create_task(self.run())

    async def stop(self):
        if self.task:
            self.task.cancel()
            with suppress(asyncio.CancelledError):
                await self.task

    async def run(self):
        while True:
            self.wakeup.clear()
            identity = self.storage.next_job()
            if identity:
                try:
                    await self.process(identity)
                except Exception:
                    log.exception("Translation job %s failed unexpectedly", identity)
                    self.storage.job_status(identity, "failed", "Internal processing error; see server log")
            else:
                await self.wakeup.wait()

    async def process(self, identity):
        job = self.storage.get_job(identity, internal=True)
        snapshot, options = job["snapshot"], job["request"]
        project = Project.model_validate({k: v for k, v in snapshot["project"].items() if k not in {"id", "revision"}})
        config = ProviderConfig.model_validate(snapshot["provider"])
        provider = self.provider_factory(config)
        self.storage.job_status(identity, "running")
        try:
            model = job["model"] or await provider.resolve_model(project.model)
            self.storage.job_status(identity, "running", model=model)
            fingerprint = hashlib.sha256(encode({"snapshot": snapshot, "model": model, "prompt_version": PROMPT_VERSION}).encode()).hexdigest()
            pending = []
            for item in job["items"]:
                if item["status"] == "completed":
                    continue
                if not item["source"].strip():
                    self.storage.item_result(identity, item["position"], "completed", item["source"])
                    continue
                if item["source"] in project.glossary:
                    # A complete dictionary match is authoritative, even when
                    # a model would mistake a short UI label for a code token.
                    translated = project.glossary[item["source"]]
                    original_tokens = list(protect(item["source"], project).tokens.values())
                    translated_tokens = list(protect(translated, project).tokens.values())
                    if not translated.strip() or original_tokens != translated_tokens:
                        self.storage.item_result(identity, item["position"], "failed", error="Exact glossary entry changes protected syntax or has empty text")
                    else:
                        self.storage.item_result(identity, item["position"], "completed", translated)
                        self.storage.cache_put(fingerprint, item["source"], translated)
                    continue
                cached = self.storage.cache_get(fingerprint, item["source"]) if options["use_cache"] else None
                if cached is not None:
                    self.storage.item_result(identity, item["position"], "completed", cached, cached=True)
                else:
                    pending.append(item)
            while pending and self.storage.job_state(identity) != "cancelled":
                batch, size = [], 0
                while pending and len(batch) < options["batch_size"]:
                    item = pending[0]
                    if batch and size + len(item["source"]) > options["max_batch_chars"]:
                        break
                    batch.append(pending.pop(0))
                    size += len(item["source"])
                await self.translate_batch(provider, model, project, job, fingerprint, batch)
            result = self.storage.get_job(identity)
            self.storage.job_status(identity, "completed_with_errors" if result["progress"]["failed"] else "completed")
        except ProviderError as exc:
            self.storage.job_status(identity, "failed", str(exc))
        finally:
            await provider.close()

    async def translate_batch(self, provider, model, project, job, fingerprint, items):
        identity = job["id"]
        remaining = items
        for attempt in range(job["request"]["retries"] + 1):
            eligible = []
            for item in remaining:
                if item["attempts"] >= job["request"]["retries"] + 1:
                    self.storage.item_result(identity, item["position"], "failed", error="Attempt budget exhausted; retry this job explicitly")
                else:
                    eligible.append(item)
            remaining = eligible
            if not remaining:
                break
            protected = {item["id"]: protect(item["source"], project) for item in remaining}
            self.storage.attempted(identity, [item["position"] for item in remaining])
            for item in remaining:
                item["attempts"] += 1
            error = "Model omitted the item or returned a duplicate ID"
            call_error = None
            try:
                output = await provider.translate(model, project, [{"id": item["id"], "text": protected[item["id"]].text} for item in remaining])
            except ProviderError as exc:
                output, error = {}, str(exc)
                call_error = str(exc)
            self.storage.record_call(identity, len(remaining), getattr(provider, "last_metrics", {}), call_error)
            failed = []
            for item in remaining:
                item_error = error
                try:
                    text = output[item["id"]]
                    if not text.strip():
                        raise ValueError("Model returned an empty translation")
                    translated = protected[item["id"]].restore(text)
                except (KeyError, ValueError) as exc:
                    if isinstance(exc, ValueError):
                        item_error = str(exc)
                    failed.append(item)
                    # Keep retryable rows pending until the attempt budget ends.
                    self.storage.item_result(identity, item["position"], "failed" if item["attempts"] >= job["request"]["retries"] + 1 else "pending", error=item_error)
                else:
                    self.storage.cache_put(fingerprint, item["source"], translated)
                    self.storage.item_result(identity, item["position"], "completed", translated)
            remaining = failed
            if remaining and attempt < job["request"]["retries"]:
                await asyncio.sleep(min(2 ** attempt, 8))