"""Managed local Supertonic TTS service.

The service owns a local Supertonic HTTP daemon. Creating/connecting the
PawFlow service starts the daemon when needed, then calls its native
``POST /v1/tts`` endpoint and returns audio bytes to ``speak`` /
``generate_audio``.
"""

import json
import logging
import os
import shlex
import shutil
import subprocess  # nosec B404 - managed local daemon uses explicit argv with shell=False.
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any, Dict

from core import ServiceFactory, ServiceError, safe_float
from core.relay_proxy_url import is_relay_proxy_url, resolve_relay_aware_url
from core.service_install import (
    assert_requirements,
    executable_requirement,
    python_venv_requirement,
    run_checked,
    write_install_state,
)
from services.base_audio_generation import BaseAudioGenerationService
from services.base_tts import BaseTTSService

logger = logging.getLogger(__name__)

_FORMATS = {"wav", "flac", "ogg"}
_CONTENT_TYPES = {
    "wav": "audio/wav",
    "flac": "audio/flac",
    "ogg": "audio/ogg",
}


def _raw_config(config: dict, key: str, default=""):
    try:
        return dict.__getitem__(config, key)
    except KeyError:
        return default


class SupertonicTTSService(BaseAudioGenerationService, BaseTTSService):
    TYPE = "supertonicTTS"
    VERSION = "1.0.0"
    NAME = "Supertonic Local TTS"
    CATEGORY = "audio"
    DESCRIPTION = (
        "Generate multilingual speech through a local Supertonic 3 HTTP "
        "daemon managed automatically by PawFlow."
    )

    def get_parameter_schema(self) -> dict:
        return {
            "base_url": {
                "type": "string", "required": False,
                "default": "http://127.0.0.1:7788",
                "description": "Supertonic server base URL. Use http://${conv.relay}/localhost:7788 for a relay-routed user endpoint.",
            },
            "allow_private_base_url": {
                "type": "boolean", "required": False, "default": False,
                "description": "Allow direct private/loopback base_url targets when auto_start is false. Prefer relay URLs for user-local endpoints.",
            },
            "auto_start": {
                "type": "boolean", "required": False, "default": True,
                "description": "Start the local Supertonic daemon automatically when the service connects.",
            },
            "auto_install": {
                "type": "boolean", "required": False, "default": True,
                "description": "Prepare a managed Supertonic runtime during service installation.",
            },
            "install_dir": {
                "type": "string", "required": False, "default": "data/runtime/supertonic3",
                "description": "Managed Supertonic runtime directory.",
            },
            "package_spec": {
                "type": "string", "required": False, "default": "supertonic[serve]>=0.1.0",
                "description": "pip package spec installed into managed runtime instead of repo_url.",
            },
            "repo_url": {
                "type": "string", "required": False, "default": "",
                "description": "Optional Supertonic repository URL to install instead of package_spec.",
            },
            "repo_ref": {
                "type": "string", "required": False, "default": "",
                "description": "Optional git ref for repo_url installs.",
            },
            "start_command": {
                "type": "string", "required": False, "default": "",
                "description": "Optional explicit command to start the Supertonic daemon.",
            },
            "startup_timeout": {
                "type": "integer", "required": False, "default": 60,
                "description": "Seconds to wait for the managed daemon to become reachable.",
            },
            "voice": {
                "type": "string", "required": False, "default": "M1",
                "description": "Built-in or imported Supertonic voice name, e.g. M1-M5 or F1-F5.",
            },
            "lang": {
                "type": "string", "required": False, "default": "na",
                "description": "Language code such as en, fr, ja, ko, or na for language-agnostic fallback.",
            },
            "steps": {
                "type": "integer", "required": False, "default": 8,
                "description": "Supertonic quality steps; higher is slower.",
            },
            "speed": {
                "type": "number", "required": False, "default": 1.05,
                "description": "Speech speed, typically 0.7 to 2.0.",
            },
            "response_format": {
                "type": "select", "required": False, "default": "wav",
                "options": sorted(_FORMATS),
                "description": "Output audio format.",
            },
            "timeout": {
                "type": "integer", "required": False, "default": 120,
                "description": "HTTP timeout in seconds.",
            },
        }

    def __init__(self, config):
        super().__init__(config)
        self.base_url = str(
            _raw_config(self.config, "base_url", "http://127.0.0.1:7788") or "http://127.0.0.1:7788").rstrip("/")
        self._raw_base_url = self.base_url
        self.allow_private_base_url = str(
            self.config.get("allow_private_base_url", False)).lower() in {"1", "true", "yes", "on"}
        self._runtime_user_id = ""
        self._runtime_conversation_id = ""
        self._runtime_agent_name = ""
        self.voice = str(self.config.get("voice") or "M1")
        self.lang = str(self.config.get("lang") or self.config.get("language") or "na")
        self.steps = int(self.config.get("steps") or self.config.get("total_steps") or 8)
        self.speed = safe_float(self.config.get("speed"), 1.05)
        self.response_format = self._format(self.config.get("response_format") or "wav")
        self.timeout = int(self.config.get("timeout") or 120)
        self.auto_start = str(self.config.get("auto_start", True)).lower() not in {"0", "false", "no"}
        self.auto_install = str(self.config.get("auto_install", True)).lower() not in {"0", "false", "no"}
        install_dir = Path(str(self.config.get("install_dir") or "data/runtime/supertonic3")).expanduser()
        if not install_dir.is_absolute():
            install_dir = Path.cwd() / install_dir
        self.install_dir = install_dir
        self.package_spec = str(self.config.get("package_spec") or "supertonic[serve]>=0.1.0")
        self.repo_url = str(self.config.get("repo_url") or "").strip()
        self.repo_ref = str(self.config.get("repo_ref") or "").strip()
        self.start_command = str(self.config.get("start_command") or "").strip()
        self.startup_timeout = int(self.config.get("startup_timeout") or 60)
        self._managed_proc = None
        self._managed_log_path = self.install_dir / "supertonic-pawflow.log"
        self._connect_lock = threading.Lock()
        self._warmup_done = False

    def set_runtime_context(self, user_id: str = "", conversation_id: str = "",
                            agent_name: str = "", **_: object):
        self._runtime_user_id = user_id or ""
        self._runtime_conversation_id = conversation_id or ""
        self._runtime_agent_name = agent_name or ""

    def _effective_base_url(self) -> str:
        return resolve_relay_aware_url(
            self._raw_base_url,
            user_id=self._runtime_user_id,
            conversation_id=self._runtime_conversation_id,
            agent_name=self._runtime_agent_name,
            allow_private=self.allow_private_base_url or self.auto_start,
            service_name="Supertonic",
            transform_relay=True,
        )

    def get_install_requirements(self):
        if self.start_command or not self.auto_install:
            return []
        reqs = [python_venv_requirement()]
        if self.repo_url:
            reqs.append(executable_requirement("git"))
        return reqs

    def prepare_install(self, reporter=None):
        if reporter:
            reporter.step("checking_requirements", "Checking Supertonic requirements")
        reqs = self.get_install_requirements()
        assert_requirements(reqs)
        if self.auto_install and not self.start_command:
            self._ensure_runtime(reporter)
        write_install_state(self.install_dir / ".pawflow_install.json", {
            "service_type": self.TYPE,
            "package_spec": self.package_spec,
            "repo_url": self.repo_url,
            "repo_ref": self.repo_ref,
        })
        return {"prepared": True, "requirements": reqs, "install_dir": str(self.install_dir)}

    def connect(self):
        """Register the service without starting the heavy local daemon.

        The managed Supertonic process is started lazily by ensure_connected(),
        which is called by the first synth request. This keeps server startup
        fast while still surfacing provider errors at first use.
        """
        parsed = self._validate_endpoint()
        self._connection = {"base_url": self.base_url, "managed": False,
                            "process": None, "lazy": True, "parsed": parsed}
        self._initialized = True
        self._log_connection("Connected lazily")

    def _validate_endpoint(self):
        resolved = resolve_relay_aware_url(
            self._raw_base_url,
            allow_private=self.allow_private_base_url or self.auto_start,
            service_name="Supertonic",
            transform_relay=False,
        )
        self.base_url = resolved
        if self.auto_start and is_relay_proxy_url(resolved):
            raise ServiceError("auto_start requires a loopback Supertonic base_url")
        parsed = urllib.parse.urlparse(self.base_url)
        if parsed.hostname not in {"127.0.0.1", "localhost", "::1"} and self.auto_start:
            raise ServiceError("auto_start requires a loopback Supertonic base_url")
        return parsed

    def _create_connection(self):
        parsed = self._validate_endpoint()
        if self._server_ready():
            return {"base_url": self.base_url, "managed": False, "process": None}
        if not self.auto_start:
            raise ServiceError(f"Supertonic server unavailable at {self.base_url}")
        proc = self._start_server(parsed)
        self._wait_ready(proc)
        self._managed_proc = proc
        return {"base_url": self.base_url, "managed": True, "process": proc}

    def _venv_python(self) -> Path:
        return self.install_dir / "venv" / ("Scripts/python.exe" if os.name == "nt" else "bin/python")

    def _resolve_start_command(self):
        if self.start_command:
            return shlex.split(self.start_command)
        venv_python = self._venv_python()
        if not venv_python.exists():
            raise ServiceError(
                "Supertonic runtime is not prepared. Install the service again "
                "or configure start_command for an existing Supertonic daemon.")
        return [
            str(venv_python),
            "-c",
            "from supertonic.cli import main; raise SystemExit(main())",
            "serve",
        ]

    def _ensure_runtime(self, reporter=None):
        python = self._venv_python()
        if python.exists():
            self._ensure_serve_runtime(python, reporter=reporter)
            if reporter:
                reporter.step("validating", "Supertonic managed runtime already exists")
            return
        self.install_dir.mkdir(parents=True, exist_ok=True)
        if reporter:
            reporter.step("creating_venv", "Creating Supertonic Python virtual environment")
        system_python = sys.executable or shutil.which("python3") or shutil.which("python")
        run_checked([system_python, "-m", "venv", str(self.install_dir / "venv")], reporter=reporter, phase="creating_venv")
        run_checked([str(python), "-m", "pip", "install", "--upgrade", "pip", "-q"], reporter=reporter, phase="installing_python_requirements")
        if self.repo_url:
            git = shutil.which("git")
            if not git:
                raise ServiceError("git is required to install Supertonic from repo_url")
            src = self.install_dir / "src" / "supertonic3"
            if not src.exists():
                if reporter:
                    reporter.step("cloning", "Cloning Supertonic repository")
                src.parent.mkdir(parents=True, exist_ok=True)
                run_checked([git, "clone", "--no-checkout", self.repo_url, str(src)], reporter=reporter, phase="cloning")
            if self.repo_ref:
                if reporter:
                    reporter.step("checkout", f"Checking out Supertonic ref {self.repo_ref}")
                run_checked([git, "-c", "safe.directory=*", "-C", str(src), "checkout", "--detach", self.repo_ref], reporter=reporter, phase="checkout")
            run_checked([str(python), "-m", "pip", "install", "-e", str(src)], reporter=reporter, phase="installing_python_requirements")
        else:
            run_checked([str(python), "-m", "pip", "install", self.package_spec], reporter=reporter, phase="installing_python_requirements")
        self._ensure_serve_runtime(python, reporter=reporter)

    def _ensure_serve_runtime(self, python: Path, reporter=None):
        """Ensure the managed Supertonic venv can run the HTTP server."""
        probe = (
            "import fastapi, uvicorn; "
            "from supertonic.cli import main"
        )
        try:
            if subprocess.run(  # nosec B603 - Python executable path is from the managed venv.
                    [str(python), "-c", probe],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                    check=False).returncode == 0:
                return
        except OSError:
            pass
        if reporter:
            reporter.step("installing_python_requirements", "Installing Supertonic server extras")
        run_checked([str(python), "-m", "pip", "install", self.package_spec], reporter=reporter, phase="installing_python_requirements")

    def ensure_connected(self):
        with self._connect_lock:
            if not self._initialized or self._connection is None:
                self.connect()
            if self._server_ready():
                return
            self._connection = self._create_connection()
            self._initialized = True

    def _close_connection(self):
        proc = self._managed_proc
        self._managed_proc = None
        if not proc or proc.poll() is not None:
            return
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait(timeout=5)

    def _server_ready(self) -> bool:
        try:
            base_url = self._effective_base_url()
        except Exception:
            return False
        probe = urllib.request.Request(
            base_url + "/v1/tts",
            data=json.dumps({
                "text": "ping",
                "voice": self.voice,
                "lang": self.lang,
                "steps": 1,
                "speed": self.speed,
                "response_format": self.response_format,
            }).encode("utf-8"),
            headers={"Content-Type": "application/json", "Accept": "audio/*"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(probe, timeout=min(5, self.timeout)) as resp:  # nosec B310 - configured local/HTTP TTS endpoint.
                return resp.status < 500
        except urllib.error.HTTPError as exc:
            return exc.code < 500
        except Exception:
            return False

    def _start_server(self, parsed) -> subprocess.Popen:
        host = parsed.hostname or "127.0.0.1"
        port = parsed.port or (443 if parsed.scheme == "https" else 7788)
        if not self.start_command and self.auto_start and self.auto_install:
            self._ensure_runtime()
        cmd = self._resolve_start_command() + ["--host", host, "--port", str(port)]
        env = os.environ.copy()
        logger.info("[SUPERTONIC] starting managed daemon: %s", " ".join(cmd))
        self._managed_log_path.parent.mkdir(parents=True, exist_ok=True)
        log_file = self._managed_log_path.open("ab")
        try:
            return subprocess.Popen(  # nosec B603 - argv is constructed locally and shell=False is the default.
                cmd,
                stdin=subprocess.DEVNULL,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                env=env,
                start_new_session=True,
            )
        except FileNotFoundError as exc:
            raise ServiceError(f"python executable not found for Supertonic: {sys.executable}") from exc

    def _wait_ready(self, proc: subprocess.Popen):
        deadline = time.time() + self.startup_timeout
        last_error = ""
        while time.time() < deadline:
            if proc.poll() is not None:
                detail = ""
                try:
                    if self._managed_log_path.exists():
                        detail = self._managed_log_path.read_text(
                            encoding="utf-8", errors="replace")[-4000:].strip()
                except Exception:
                    detail = ""
                raise ServiceError(f"Supertonic daemon exited during startup: {detail}")
            if self._server_ready():
                return
            time.sleep(0.5)
        proc.terminate()
        try:
            proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            proc.kill()
        raise ServiceError(
            f"Supertonic daemon did not become ready at {self.base_url} within {self.startup_timeout}s{(': ' + last_error) if last_error else ''}")

    @staticmethod
    def _format(value: str) -> str:
        fmt = str(value or "wav").lower().strip()
        if fmt not in _FORMATS:
            raise ServiceError(
                f"unsupported Supertonic response_format {fmt!r}; expected one of {sorted(_FORMATS)}")
        return fmt

    def _post_tts(self, body: Dict[str, Any]) -> tuple[bytes, str]:
        url = f"{self._effective_base_url()}/v1/tts"
        payload = json.dumps(body).encode("utf-8")
        req = urllib.request.Request(
            url,
            data=payload,
            headers={
                "Content-Type": "application/json",
                "Accept": "audio/*",
                "User-Agent": "PawFlow-Agent/1.0",
            },
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:  # nosec B310 - configured local/HTTP TTS endpoint.
                content_type = resp.headers.get(
                    "Content-Type", _CONTENT_TYPES.get(body.get("response_format", "wav"), "audio/wav"))
                return resp.read(), content_type
        except urllib.error.HTTPError as exc:
            detail = exc.read()[:500].decode("utf-8", errors="replace")
            raise ServiceError(
                f"Supertonic API error POST /v1/tts ({exc.code}): {detail}") from exc
        except urllib.error.URLError as exc:
            raise ServiceError(f"Supertonic server unavailable at {url}: {exc}") from exc

    def generate(self, prompt: str = "", text: str = "", voice: str = "",
                 lang: str = "", language: str = "", steps: int = 0,
                 total_steps: int = 0, speed: float = 0,
                 response_format: str = "", **kwargs) -> dict:
        return self.text_to_speech(
            text=text or prompt,
            voice=voice,
            lang=lang or language,
            steps=steps or total_steps,
            speed=speed,
            response_format=response_format,
            **kwargs,
        )

    def text_to_speech(self, text: str = "", voice: str = "", lang: str = "",
                       language: str = "", steps: int = 0,
                       total_steps: int = 0, speed: float = 0,
                       response_format: str = "", **kwargs) -> dict:
        if not text:
            raise ServiceError("No text provided")
        self.ensure_connected()

        fmt = self._format(response_format or self.response_format)
        body: Dict[str, Any] = {
            "text": text,
            "voice": voice or self.voice,
            "lang": lang or language or self.lang,
            "steps": int(steps or total_steps or self.steps),
            "speed": safe_float(speed, self.speed) if speed else self.speed,
            "response_format": fmt,
        }
        for key, value in kwargs.items():
            if key in body or key in {
                "destination", "path", "service", "audio_service", "_service",
                "model", "prompt", "lyrics", "instrumental", "style",
            }:
                continue
            if value is not None and value != "":
                body[key] = value

        audio_bytes, content_type = self._post_tts(body)
        if not audio_bytes:
            raise ServiceError("Supertonic returned empty audio")
        logger.info("[SUPERTONIC] tts ok: %d bytes (%s)", len(audio_bytes), content_type)
        return {
            "audio_bytes": audio_bytes,
            "content_type": content_type or _CONTENT_TYPES[fmt],
            "source_url": "",
        }

    def warmup(self, voice: str = "", language: str = "", **_kwargs) -> None:
        if self._warmup_done:
            return
        self.ensure_connected()
        body: Dict[str, Any] = {
            "text": "OK",
            "voice": voice or self.voice,
            "lang": language or self.lang,
            "steps": self.steps,
            "speed": self.speed,
            "response_format": "wav",
        }
        audio_bytes, _content_type = self._post_tts(body)
        if not audio_bytes:
            raise ServiceError("Supertonic returned empty warmup audio")
        self._warmup_done = True

    def speak(self, text: str, voice: str = "", language: str = "",
              **kwargs) -> dict:
        return self.text_to_speech(
            text=text,
            voice=voice,
            language=language,
            **kwargs,
        )


ServiceFactory.register(SupertonicTTSService)
