"""Orchestrate CPU preparation, the synchronous GPU call, and a single CPU export."""
from contextlib import contextmanager
import html
import json
import logging
from pathlib import Path
import threading
import time

import gradio as gr

from audio_files import export_take, prune_outputs, read_reference
from lifecycle import AlreadyRunning, Cancelled, Jobs, check_cancelled
from requests_schema import InputError, validate

log = logging.getLogger("zonos.controller")


def message(text, error=False, busy=False):
    return f'<p id="status-state" data-busy="{str(busy).lower()}" role="{"alert" if error else "status"}">{html.escape(text)}</p>'


class Controller:
    def __init__(self, engine, runtime_dir, voice_names, aligner=None):
        self.engine = engine
        self.aligner = aligner
        self.root = Path(runtime_dir)
        self.outputs = self.root / "outputs"
        self.outputs.mkdir(parents=True, exist_ok=True)
        self.jobs = Jobs(self.root / "jobs")
        self.gate = threading.Semaphore(1)
        self.voice_names = set(voice_names)

    @contextmanager
    def inference_slot(self, job):
        acquired = False
        try:
            while not acquired:
                check_cancelled(job.marker)
                acquired = self.gate.acquire(timeout=.1)
            check_cancelled(job.marker)
            yield
        finally:
            if acquired:
                self.gate.release()

    def cancel(self, request: gr.Request):
        owner = request.session_hash if request is not None else None
        if not owner:
            return gr.skip(), gr.skip()
        if self.jobs.cancel(owner):
            return message("Cancellation requested. Waiting for the current operation to stop.", busy=True), gr.update(interactive=False)
        return gr.skip(), gr.skip()

    def on_disconnect(self, request: gr.Request):
        if request is not None and request.session_hash:
            self.jobs.cancel(request.session_hash)

    def generate(self, text, picker, source, clip, consent, normalize, seed, randomize,
                 temperature, top_p, top_k, repetition_penalty, rate, request: gr.Request,
                 progress=gr.Progress(track_tqdm=False)):
        # Outputs: waveform/player data, status, main button, cancel button, seed.
        job = None
        final = None
        pending_error = None
        exported = None
        keep_export = False
        try:
            owner = request.session_hash if request is not None else None
            if not owner:
                raise InputError("Missing session. Reload the page or use gradio_client.Client.")
            request = validate(text, picker, source, seed, randomize, normalize, temperature,
                               top_p, top_k, repetition_penalty, rate, self.voice_names,
                               bool(clip), consent)
            job = self.jobs.begin(owner)
            yield (gr.skip(), message("Request received.", busy=True), gr.update(interactive=False, value="Working…"),
                   gr.update(visible=True, interactive=True), gr.skip())
            progress(None, desc="Preparing text")
            reference = read_reference(clip) if request.source == "Voice clip" else None
            check_cancelled(job.marker)
            prepared = self.engine.prepare(request, reference)
            del reference
            progress(None, desc="Waiting for an inference slot")
            with self.inference_slot(job):
                progress(None, desc="Requesting GPU")
                outcome = self.engine.run(prepared, job.marker, progress)
                # The decorated call has returned; its SDK release path completed.
                log.info("gpu_call_returned request=%s", job.id)
            del prepared
            check_cancelled(job.marker)
            if outcome.cancelled:
                raise Cancelled()
            yield (gr.skip(), message("Finalizing audio…", busy=True), gr.update(interactive=False, value="Finalizing…"),
                   gr.update(visible=True, interactive=False), gr.skip())
            progress(None, desc="Finalizing audio")
            check_cancelled(job.marker)
            exported, result = export_take(outcome.pcm, self.outputs, seed=request.seed,
                                 request_seconds=time.monotonic()-job.started,
                                 inference_seconds=outcome.seconds)
            result.update(text=request.text, words=[], alignment="unavailable")
            if self.aligner is not None:
                yield (gr.skip(), message("Synchronizing words…", busy=True), gr.update(interactive=False, value="Synchronizing…"),
                       gr.update(visible=True, interactive=True), gr.skip())
                try:
                    from alignment import attach_timings
                    progress(None, desc="Preparing word synchronization")
                    alignment_inputs = self.aligner.prepare(outcome.pcm, request.text)
                    check_cancelled(job.marker)
                    with self.inference_slot(job):
                        progress(None, desc="Requesting GPU for word synchronization")
                        words = self.aligner.run(alignment_inputs, job.marker, progress)
                        log.info("alignment_gpu_returned request=%s", job.id)
                    check_cancelled(job.marker)
                    result["words"] = attach_timings(request.text, words, result["duration"])
                    result["alignment"] = "ready"
                except Cancelled:
                    raise
                except Exception as exc:
                    check_cancelled(job.marker)
                    log.warning("alignment_unavailable request=%s error_type=%s", job.id, type(exc).__name__)
                finally:
                    alignment_inputs = None
            del outcome
            check_cancelled(job.marker)
            prune_outputs(self.outputs)
            result["request_seconds"] = round(time.monotonic()-job.started, 2)
            final = (json.dumps(result), message("Ready to play." if result["alignment"] == "ready" or self.aligner is None else "Audio ready. Word synchronization unavailable for this take."),
                     gr.update(interactive=True, value="Generate speech"),
                     gr.update(visible=False, interactive=True), request.seed)
            keep_export = True
        except AlreadyRunning:
            # Do not re-enable controls belonging to the already-running request.
            yield gr.skip(), message("A request is already running in this session.", busy=True), gr.skip(), gr.skip(), gr.skip()
            raise gr.Error("A request is already running in this session.")
        except Cancelled:
            final = (gr.skip(), message("Cancelled. Your previous audio is still available."),
                     gr.update(interactive=True, value="Generate speech"),
                     gr.update(visible=False, interactive=True), gr.skip())
        except InputError as exc:
            pending_error = gr.Error(str(exc))
            final = (gr.skip(), message(str(exc), error=True), gr.update(interactive=True, value="Generate speech"),
                     gr.update(visible=False, interactive=True), gr.skip())
        except Exception as exc:
            # Do not echo exception bodies that may include user text, audio or paths.
            log.error("request_failed request=%s error_type=%s", job.id if job else "validation", type(exc).__name__)
            text = "Generation failed. Your previous audio is still available. Try again or shorten the text."
            if job and job.marker.exists():
                text = "Cancelled. Your previous audio is still available."
            elif isinstance(exc, gr.Error):
                # ZeroGPU owns quota/authentication errors and its native modal.
                text = "The GPU request could not complete. Check the Hugging Face notice, then try again."
            if not (job and job.marker.exists()):
                pending_error = exc if isinstance(exc, gr.Error) else gr.Error(text)
            final = (gr.skip(), message(text, error=True), gr.update(interactive=True, value="Generate speech"),
                     gr.update(visible=False, interactive=True), gr.skip())
        finally:
            if exported is not None and not keep_export:
                exported.unlink(missing_ok=True)
            if job is not None:
                self.jobs.finish(job)
        # Finishing happens BEFORE the final yield, even if the consumer never
        # resumes or explicitly closes the generator after receiving the result.
        if final is not None:
            yield final
        if pending_error is not None:
            # Preserve the native error notice and an API failure, after controls
            # have been restored and the previous audio has been preserved.
            raise pending_error