"""Nonblocking, cancellable requests to a resident Chatterbox worker in WSL/Linux."""
import hashlib
import json
import os
from pathlib import Path
import queue
import shlex
import subprocess
import threading

ROOT = Path(__file__).resolve().parent
REFERENCE = ROOT / ".speech/voices/teddy-chatterbox-reference.wav"


def linux_path(path):
    path = Path(path).resolve()
    if os.name != "nt":
        return str(path)
    return "/mnt/" + path.drive[0].lower() + path.as_posix()[2:]


def voice_id():
    return "teddy-chatterbox-turbo-v1-" + hashlib.sha256(REFERENCE.read_bytes()).hexdigest()[:16]


class Job:
    def __init__(self, text, output):
        self.text, self.output = text, output
        self.cancelled = threading.Event()
        self.returncode = None
    def poll(self):
        return self.returncode
    def terminate(self):
        # In-flight GPU work may finish and populate cache, but must never play.
        self.cancelled.set()
    def wait(self, timeout=None):
        return self.returncode


class Engine:
    def __init__(self):
        self.jobs = queue.Queue()
        self.ready = threading.Event()
        self.error = None
        self.process = None
        self.thread = threading.Thread(target=self._run, daemon=True)
        self.thread.start()
    def submit(self, text, output):
        job = Job(text, output)
        self.jobs.put(job)
        return job
    def close(self):
        self.jobs.put(None)
    def _run(self):
        current = None
        try:
            home = os.environ.get("TEDDY_CHATTERBOX_HOME", "~/teddy-chatterbox")
            # Expand the conventional home prefix in the Linux shell only.
            home_arg = '"$HOME"' + shlex.quote(home[1:]) if home.startswith("~/") else shlex.quote(home)
            script = linux_path(ROOT / "chatterbox_worker.py")
            command = (f"export HF_HOME={home_arg}/models; export HF_HUB_DISABLE_XET=1; "
                       f"exec {home_arg}/venv/bin/python -u {shlex.quote(script)} {shlex.quote(linux_path(REFERENCE))}")
            args = (["wsl", "-d", "Ubuntu", "--", "bash", "-lc", command] if os.name == "nt"
                    else ["bash", "-lc", command])
            log = ROOT / ".speech/chatterbox.log"
            log.parent.mkdir(parents=True, exist_ok=True)
            with log.open("a", encoding="utf-8") as logs:
                self.process = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                    stderr=logs, text=True, encoding="utf-8",
                    creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0)
                if not json.loads(self.process.stdout.readline()).get("ready"):
                    raise RuntimeError("Chatterbox did not become ready")
                self.ready.set()
                while True:
                    current = self.jobs.get()
                    if current is None:
                        break
                    if current.cancelled.is_set():
                        current.returncode = 1
                        continue
                    self.process.stdin.write(json.dumps({"text": current.text,
                        "output": linux_path(current.output)}) + "\n")
                    self.process.stdin.flush()
                    result = json.loads(self.process.stdout.readline())
                    current.returncode = 0 if result.get("ok") else 1
                    if not result.get("ok"):
                        print("Chatterbox synthesis: " + result.get("error", "unknown error"), flush=True)
                    current = None
        except Exception as error:
            self.error = str(error)
            if current is not None:
                current.returncode = 1
            print("Chatterbox worker failed: " + self.error, flush=True)
        finally:
            if self.process:
                self.process.stdin.close()
                try:
                    self.process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    self.process.terminate()