"""Agent-facing handlers for extra media capabilities.

Each handler is provider-agnostic: it calls a resolver `() -> (service,
error_msg)` wired by `AgentToolConfigMixin` at request time, then
delegates to the service's public method. File storage is routed
through `core.storage_resolver.StorageResolver`, same pattern as
`ImageGenerationHandler` / `VideoGenerationHandler`.
"""

import functools
import logging
import os
import time
from typing import Any, Dict

from core.tool_handler import ToolHandler

logger = logging.getLogger(__name__)

_SERVICE_ARG_NAMES = (
    "service", "image_service", "video_service", "audio_service", "voice_service")
_INTERNAL_TTS_ARG_NAMES = ("_tts_storage_ttl", "transient", "transient_ttl", "ttl")


def _resolve_filestore_url(url: str, base_url: str, service=None) -> str:
    """Convert fs://filestore/<id>/<name> to HTTP unless service reads it locally."""
    if not url or not url.startswith("fs://filestore/"):
        return url
    if service is not None and getattr(service, "ACCEPTS_FILESTORE_URLS", False):
        return url
    rest = url[len("fs://filestore/"):]
    fid = rest.split("/", 1)[0]
    if not fid:
        return url
    return f"{base_url.rstrip('/')}/files/{fid}"


def _write_result(user_id: str, conversation_id: str,
                   destination: str, filename: str,
                   payload: bytes, content_type: str) -> Dict[str, Any]:
    """Persist the generated bytes via StorageResolver, return its dict."""
    from core.storage_resolver import StorageResolver
    resolver = StorageResolver(
        user_id=user_id, conversation_id=conversation_id or "")
    return resolver.write(destination, filename, payload, content_type)


def _write_result_file(user_id: str, conversation_id: str,
                       destination: str, filename: str,
                       source_path: str, content_type: str) -> Dict[str, Any]:
    """Persist a generated local file via StorageResolver."""
    from core.storage_resolver import StorageResolver
    resolver = StorageResolver(
        user_id=user_id, conversation_id=conversation_id or "")
    return resolver.write_file(destination, filename, source_path, content_type)


def _provider_identity(svc) -> str:
    runtime = getattr(svc, "config", {}).get("package_runtime", {}) \
        if isinstance(getattr(svc, "config", None), dict) else {}
    if runtime:
        package = str(runtime.get("package") or "")
        object_id = str(runtime.get("object_id") or "")
        if package and object_id:
            return f"pfp:{package}/{object_id}"
    return getattr(svc, "TYPE", svc.__class__.__name__)


def _provider_version(svc) -> str:
    config = getattr(svc, "config", {})
    runtime = config.get("package_runtime", {}) if isinstance(config, dict) else {}
    if runtime:
        installed = config.get("installed_from", {}) if isinstance(config, dict) else {}
        version = str(runtime.get("version") or "")
        digest = str(installed.get("hash") or "")
        return f"{version}:{digest}" if digest else version
    return getattr(svc, "VERSION", "0")


def _voice_id_value(value) -> str:
    """Normalize provider voice-id results to the provider-opaque string."""
    if isinstance(value, dict):
        return str(value.get("voice_id") or value.get("id") or "")
    return str(value or "")


class _CapabilityHandlerBase(ToolHandler):
    """Shared wiring for the capability handlers (3D/upscale/try-on/...)."""

    _base_url: str = "http://localhost:9090"
    _service_resolver = None  # () -> (service, error_msg)
    _user_id: str = ""
    _conversation_id: str = ""
    _agent_name: str = ""
    _share = None  # active TemporaryPublicRefs for the current execute() call

    def __init_subclass__(cls, **kwargs):
        """Wrap each concrete handler's execute() in a temporary-share scope.

        Every capability handler routes its reference inputs through
        ``_rewrite``; wrapping ``execute`` here means each call shares the
        referenced FileStore files as public, gateway-key URLs for the
        duration of the call and revokes them afterwards — success or
        failure — without per-handler boilerplate.
        """
        super().__init_subclass__(**kwargs)
        raw_execute = cls.__dict__.get("execute")
        if raw_execute is None:
            return  # inherits an already-wrapped execute

        @functools.wraps(raw_execute)
        def _execute_with_share(self, arguments, _raw=raw_execute):
            from core.media_share import TemporaryPublicRefs
            self._share = TemporaryPublicRefs(self._base_url, self._user_id)
            try:
                return _raw(self, arguments)
            finally:
                self._share.restore()
                self._share = None

        cls.execute = _execute_with_share

    def set_base_url(self, base_url: str):
        self._base_url = base_url.rstrip("/")

    def set_user_id(self, user_id: str):
        self._user_id = user_id

    def set_conversation_id(self, conversation_id: str):
        self._conversation_id = conversation_id

    def set_agent_name(self, agent_name: str):
        self._agent_name = agent_name

    def set_service_resolver(self, resolver):
        self._service_resolver = resolver

    def _get_service(self, arguments: Dict[str, Any] = None):
        arguments = arguments or {}
        service_name = (arguments.get("service")
                        or arguments.get("image_service")
                        or arguments.get("video_service")
                        or arguments.get("audio_service")
                        or arguments.get("voice_service")
                        or "")
        if service_name:
            try:
                from core.service_registry import ServiceRegistry
                svc = ServiceRegistry.get_instance().resolve(
                    service_name,
                    user_id=self._user_id,
                    conv_id=self._conversation_id,
                )
            except Exception as exc:
                return None, f"media service '{service_name}' failed to resolve: {exc}"
            if not svc:
                return None, f"media service '{service_name}' not found or not connected"
            if hasattr(svc, "set_runtime_context"):
                svc.set_runtime_context(
                    user_id=self._user_id,
                    conversation_id=self._conversation_id,
                    agent_name=self._agent_name,
                )
            if hasattr(svc, "set_callback_base_url"):
                svc.set_callback_base_url(self._base_url)
            return svc, ""
        if not self._service_resolver:
            return None, "no service resolver configured"
        svc, err = self._service_resolver()
        if svc and hasattr(svc, "set_runtime_context"):
            svc.set_runtime_context(
                user_id=self._user_id,
                conversation_id=self._conversation_id,
                agent_name=self._agent_name,
            )
        if svc and hasattr(svc, "set_callback_base_url"):
            svc.set_callback_base_url(self._base_url)
        return svc, err

    def _rewrite(self, url: str, service=None) -> str:
        # Inside execute() the call runs within a temporary-share scope, so
        # FileStore refs become public gateway-key URLs that are revoked when
        # the call returns. Outside that scope, fall back to the legacy form.
        share = getattr(self, "_share", None)
        if share is not None:
            return share.public_url(url, service=service)
        return _resolve_filestore_url(url, self._base_url, service=service)

    def _persist(self, destination: str, filename: str, result: Dict[str, Any],
                 label: str) -> str:
        """Write the service result to storage + return a user-facing message."""
        ct = result.get("content_type", "application/octet-stream")
        payload = result.get("bytes") or result.get("image_bytes") \
            or result.get("video_bytes") or result.get("audio_bytes") or b""
        source_path = result.get("path") or result.get("image_path") \
            or result.get("video_path") or result.get("audio_path") or ""
        if source_path:
            try:
                write_result = _write_result_file(
                    self._user_id, self._conversation_id,
                    destination, filename, str(source_path), ct)
            finally:
                if result.get("_delete_media_path"):
                    try:
                        os.unlink(str(source_path))
                    except OSError:
                        pass
            if write_result.get("file_id"):
                url = f"fs://filestore/{write_result['file_id']}/{filename}"
                return f"{label}: {url}\nfile_id: {write_result['file_id']}"
            return (f"{label} saved to "
                    f"{write_result.get('destination', destination)}: "
                    f"{write_result.get('path', filename)}")
        if not payload:
            return f"Error: {label} returned no payload"
        write_result = _write_result(
            self._user_id, self._conversation_id,
            destination, filename, payload, ct)
        if write_result.get("file_id"):
            url = f"fs://filestore/{write_result['file_id']}/{filename}"
            return f"{label}: {url}\nfile_id: {write_result['file_id']}"
        return (f"{label} saved to "
                f"{write_result.get('destination', destination)}: "
                f"{write_result.get('path', filename)}")


# ── 3D ─────────────────────────────────────────────────────────────────


class Generate3DHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "generate_3d"

    @property
    def description(self) -> str:
        return (
            "Generate a 3D model (GLB / GLTF / OBJ / USDZ) from an "
            "input image or a text prompt. Uses the active 3D service "
            "(Hunyuan3D, Rodin, Trellis, Tripo3D, …). Provide `image_url` "
            "for image-to-3D (preferred, higher quality) or `prompt` for "
            "text-to-3D. Saves the result to FileStore (default) or a "
            "filesystem service when `destination` + `path` are given."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "Optional text prompt for text-to-3D"},
                "image_url": {"type": "string", "description": "Source image URL for image-to-3D (HTTP or fs://filestore/<id>/<name>)"},
                "destination": {"type": "string", "description": "'filestore' (default) or relay service name"},
                "path": {"type": "string", "description": "Filename when saving to a filesystem service"},
                "model": {"type": "string", "description": "Override the active 3D model (e.g. 'hyper3d-rodin-259', 'tripo3d-v2-5-413', 'hunyuan3d-3-0-api-294')."},
            },
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no 3D generation service available'}"
        prompt = arguments.get("prompt", "") or ""
        image_url = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        if not prompt and not image_url:
            return "Error: provide `prompt` or `image_url`"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path", "prompt", "image_url", *_SERVICE_ARG_NAMES)}
            r = svc.generate_3d(prompt=prompt, image_url=image_url, **kwargs)
            ext = {"model/gltf-binary": "glb", "model/gltf+json": "gltf",
                   "model/obj": "obj", "model/vnd.usdz+zip": "usdz"}.get(
                       r.get("content_type", "").split(";")[0].strip(), "glb")
            filename = arguments.get("path") or (
                f"model_{int(time.time())}.{ext}")
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "3D model generated")
        except Exception as e:
            return f"Error generating 3D model: {e}"


# ── Upscale ────────────────────────────────────────────────────────────


class UpscaleImageHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "upscale_image"

    @property
    def description(self) -> str:
        return (
            "Upscale an image (2x / 4x / 8x) using an AI upscaler "
            "(SeedVR, Crystal Upscaler, Topaz, …). Pass the source via "
            "`image_url` (HTTP or fs://filestore/<id>/<name>). Saves the "
            "result to FileStore (default) or a filesystem service."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "image_url": {"type": "string", "description": "Source image URL"},
                "scale": {"type": "integer", "description": "Upscale factor (2, 4 — default 2)"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override the upscale model (e.g. 'seedvr-upscale', 'crystal-upscaler', 'topaz-upscale-video-753', 'bria-rmbg-2-0-682')."},
            },
            "required": ["image_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no upscale service available'}"
        image_url = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        if not image_url:
            return "Error: `image_url` is required"
        scale = int(arguments.get("scale", 2))
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path", "image_url", "scale", *_SERVICE_ARG_NAMES)}
            r = svc.upscale(image_url=image_url, scale=scale, **kwargs)
            ct = r.get("content_type", "image/png").split(";")[0].strip()
            ext = {"image/png": "png", "image/jpeg": "jpg",
                   "image/webp": "webp"}.get(ct, "png")
            filename = arguments.get("path") or (
                f"upscaled_{int(time.time())}.{ext}")
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Upscaled image")
        except Exception as e:
            return f"Error upscaling image: {e}"


# ── Upscale Video ─────────────────────────────────────────────────────


class UpscaleVideoHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "upscale_video"

    @property
    def description(self) -> str:
        return (
            "Upscale a video (2x / 4x) using an AI upscaler "
            "(SeedVR Video, Topaz, ...). Pass the source via "
            "`video_url` (HTTP or fs://filestore/<id>/<name>). Saves the "
            "result to FileStore (default) or a filesystem service."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "video_url": {"type": "string", "description": "Source video URL"},
                "scale": {"type": "integer", "description": "Upscale factor (2, 4 — default 2)"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override the upscale model (e.g. 'seedvr-upscale-video', 'topaz-upscale-video-753')."},
            },
            "required": ["video_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no upscale service available'}"
        video_url = self._rewrite(arguments.get("video_url", "") or "", service=svc)
        if not video_url:
            return "Error: `video_url` is required"
        if not hasattr(svc, 'upscale_video'):
            return "Error: the active upscale service does not support video upscaling"
        scale = int(arguments.get("scale", 2))
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path", "video_url", "scale", *_SERVICE_ARG_NAMES)}
            r = svc.upscale_video(video_url=video_url, scale=scale, **kwargs)
            filename = arguments.get("path") or (
                f"upscaled_{int(time.time())}.mp4")
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Upscaled video")
        except Exception as e:
            return f"Error upscaling video: {e}"


# ── Describe Image ───────────────────────────────────────────────────


class DescribeImageHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "describe_image"

    @property
    def description(self) -> str:
        return (
            "Describe the content of an image using an AI model "
            "(Ideogram v2 / Turbo). Returns a text description of the "
            "image content, style, and composition. Pass the source via "
            "`image_url` (HTTP or fs://filestore/<id>/<name>)."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "image_url": {"type": "string", "description": "Source image URL to describe"},
                "model": {"type": "string", "description": "Override the model (e.g. 'ideogram', 'ideogram-turbo')."},
            },
            "required": ["image_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no image service available'}"
        image_url = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        if not image_url:
            return "Error: `image_url` is required"
        if not hasattr(svc, 'describe_image'):
            return "Error: the active image service does not support describe_image"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("image_url", *_SERVICE_ARG_NAMES)}
            r = svc.describe_image(image_url=image_url, **kwargs)
            return f"Image description: {r.get('description', '(no description)')}"
        except Exception as e:
            return f"Error describing image: {e}"


# ── Remix Image ─────────────────────────────────────────────────────


class RemixImageHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "remix_image"

    @property
    def description(self) -> str:
        return (
            "Remix an image: generate a new image inspired by a source image "
            "and a text prompt. The output blends the style/content of the "
            "source with the prompt direction. Uses Ideogram v2 / Turbo remix. "
            "Pass the source via `image_url` (HTTP or fs://filestore/<id>/<name>)."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "Text prompt describing the desired remix"},
                "image_url": {"type": "string", "description": "Source image URL to remix"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override the model (e.g. 'ideogram', 'ideogram-turbo')."},
            },
            "required": ["prompt", "image_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no image service available'}"
        image_url = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        prompt = arguments.get("prompt", "")
        if not image_url or not prompt:
            return "Error: `prompt` and `image_url` are required"
        if not hasattr(svc, 'remix_image'):
            return "Error: the active image service does not support remix_image"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path", "image_url", "prompt", *_SERVICE_ARG_NAMES)}
            r = svc.remix_image(prompt=prompt, image_url=image_url, **kwargs)
            ct = r.get("content_type", "image/png").split(";")[0].strip()
            ext = {"image/png": "png", "image/jpeg": "jpg", "image/webp": "webp"}.get(ct, "png")
            filename = arguments.get("path") or f"remix_{int(time.time())}.{ext}"
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Remixed image")
        except Exception as e:
            return f"Error remixing image: {e}"


# ── Remove Background ─────────────────────────────────────────────────


class RemoveBackgroundHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "remove_background"

    @property
    def description(self) -> str:
        return (
            "Remove the background from an image using an AI model "
            "(Bria RMBG 2.0, VEED, ...). Pass the source via "
            "`image_url` (HTTP or fs://filestore/<id>/<name>). Returns "
            "a PNG with transparent background."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "image_url": {"type": "string", "description": "Source image URL"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override the model (e.g. 'bria-rmbg-2-0-682')."},
            },
            "required": ["image_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no upscale/background service available'}"
        image_url = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        if not image_url:
            return "Error: `image_url` is required"
        if not hasattr(svc, 'remove_background'):
            return "Error: the active service does not support remove_background"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path", "image_url", *_SERVICE_ARG_NAMES)}
            r = svc.remove_background(image_url=image_url, **kwargs)
            filename = arguments.get("path") or (
                f"nobg_{int(time.time())}.png")
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Background removed")
        except Exception as e:
            return f"Error removing background: {e}"


# ── Try-on ─────────────────────────────────────────────────────────────


class TryOnHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "try_on"

    @property
    def description(self) -> str:
        return (
            "Virtual try-on: dress a person image with a garment image. "
            "Produces a single output image where the person wears the "
            "garment. Uses the active try-on service (Kling VTON, Fashn, "
            "IDM-VTON, …). Both inputs are URLs (HTTP or fs://filestore/...)."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "person_image": {"type": "string", "description": "URL of the person image"},
                "garment_image": {"type": "string", "description": "URL of the garment image"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override the try-on model (e.g. 'fashn-virtual-try-on', 'idm-vton-api', 'kling-ai-vton', 'glass-virtual-try-on')."},
            },
            "required": ["person_image", "garment_image"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no try-on service available'}"
        person = self._rewrite(arguments.get("person_image", "") or "", service=svc)
        garment = self._rewrite(arguments.get("garment_image", "") or "", service=svc)
        if not person or not garment:
            return "Error: `person_image` and `garment_image` are required"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path",
                                   "person_image", "garment_image", *_SERVICE_ARG_NAMES)}
            r = svc.try_on(person_image=person, garment_image=garment,
                           **kwargs)
            ct = r.get("content_type", "image/png").split(";")[0].strip()
            ext = {"image/png": "png", "image/jpeg": "jpg"}.get(ct, "png")
            filename = arguments.get("path") or f"tryon_{int(time.time())}.{ext}"
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Try-on result")
        except Exception as e:
            return f"Error running try-on: {e}"


# ── Lipsync ────────────────────────────────────────────────────────────


class LipsyncHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "lipsync"

    @property
    def description(self) -> str:
        return (
            "Generate a lipsync video: drive a face (video or still image) "
            "with an audio track. Produces an MP4 where the face speaks the "
            "audio. Uses the active lipsync service (OmniHuman, Kling "
            "Avatar, Sync Lipsync, …). Provide `audio_url` and either "
            "`video_url` (preferred) or `image_url`."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "video_url": {"type": "string", "description": "Source face video (URL)"},
                "image_url": {"type": "string", "description": "Source face image (URL) — used if `video_url` is absent"},
                "audio_url": {"type": "string", "description": "Audio track (URL) to drive the face"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override the lipsync model (e.g. 'omnihuman', 'bytedance-omnihuman-v1-5-290', 'kling-ai-avatar-v2-pro-789')."},
            },
            "required": ["audio_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no lipsync service available'}"
        video = self._rewrite(arguments.get("video_url", "") or "", service=svc)
        image = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        audio = self._rewrite(arguments.get("audio_url", "") or "", service=svc)
        if not audio or not (video or image):
            return ("Error: `audio_url` is required plus either `video_url` "
                    "or `image_url`")
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path",
                                   "video_url", "image_url", "audio_url", *_SERVICE_ARG_NAMES)}
            r = svc.lipsync(video_url=video, image_url=image,
                            audio_url=audio, **kwargs)
            filename = arguments.get("path") or f"lipsync_{int(time.time())}.mp4"
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Lipsync video")
        except Exception as e:
            return f"Error running lipsync: {e}"


# ── Trainer ────────────────────────────────────────────────────────────


class TrainImageModelHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "train_image_model"

    @property
    def description(self) -> str:
        return (
            "Fine-tune an image model on a dataset (LoRA training on Flux "
            "variants, or full training on supported bases). Returns a "
            "URL to the trained LoRA / checkpoint. This is an async job — "
            "the service polls internally and the call waits for completion "
            "(no timeout, cancel via /stop). Pass `dataset_url` (zip archive "
            "of training images) and an optional `base_model` name."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "dataset_url": {"type": "string", "description": "URL (or fs://) to the training dataset (zip)"},
                "base_model": {"type": "string", "description": "Base model to fine-tune (optional)"},
                "steps": {"type": "integer", "description": "Training steps (provider-dependent default)"},
                "learning_rate": {"type": "number"},
                "trigger_word": {"type": "string"},
                "model": {"type": "string", "description": "Override the trainer model (e.g. 'flux-2-pro-text-to-image-trainer-712', 'flux-2-pro-image-to-image-trainer-831', 'qwen-image-edit-plus-trainer')."},
            },
            "required": ["dataset_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no trainer service available'}"
        dataset = self._rewrite(arguments.get("dataset_url", "") or "", service=svc)
        base_model = arguments.get("base_model", "") or ""
        if not dataset:
            return "Error: `dataset_url` is required"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("dataset_url", "base_model", *_SERVICE_ARG_NAMES)}
            r = svc.train(dataset_url=dataset, base_model=base_model,
                          **kwargs)
            lora = r.get("lora_url") or r.get("source_url") or ""
            if not lora:
                return f"Error: training completed but no LoRA URL returned"
            return f"Training complete. LoRA: {lora}"
        except Exception as e:
            return f"Error training image model: {e}"


# ── Speech to Video ─────────────────────────────────────────────────


class SpeechToVideoHandler(_CapabilityHandlerBase):
    @property
    def name(self) -> str:
        return "speech_to_video"

    @property
    def description(self) -> str:
        return (
            "Generate a lip-synced video from a face image and audio track "
            "(Wan 2.2 Speech-to-Video). The output video shows the person "
            "speaking/singing along with the audio. Requires `image_url` "
            "(face photo) and `audio_url` (speech/music)."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "Style/scene description (e.g. 'beach vacation, man in sunglasses')"},
                "image_url": {"type": "string", "description": "Face image URL (HTTP or fs://filestore/<id>/<name>)"},
                "audio_url": {"type": "string", "description": "Audio track URL (speech or music)"},
                "resolution": {"type": "string", "description": "Output resolution: '480p', '580p', '720p' (default '480p')"},
                "destination": {"type": "string"},
                "path": {"type": "string"},
                "model": {"type": "string", "description": "Override model (default 'wan2.2-s2v')."},
            },
            "required": ["image_url", "audio_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no video service available'}"
        image_url = self._rewrite(arguments.get("image_url", "") or "", service=svc)
        audio_url = self._rewrite(arguments.get("audio_url", "") or "", service=svc)
        if not image_url or not audio_url:
            return "Error: `image_url` and `audio_url` are required"
        if not hasattr(svc, 'speech_to_video'):
            return "Error: the active video service does not support speech_to_video"
        try:
            kwargs = {k: v for k, v in arguments.items()
                      if k not in ("destination", "path", "image_url", "audio_url", *_SERVICE_ARG_NAMES)}
            r = svc.speech_to_video(image_url=image_url, audio_url=audio_url, **kwargs)
            filename = arguments.get("path") or f"s2v_{int(time.time())}.mp4"
            destination = arguments.get("destination", "filestore")
            return self._persist(destination, filename, r, "Speech-to-video generated")
        except Exception as e:
            return f"Error generating speech-to-video: {e}"


# ── Voice Clone ─────────────────────────────────────────────────────


class CloneVoiceHandler(_CapabilityHandlerBase):
    """Create (or reuse) a voice clone resource from a reference audio sample.

    The resulting voice is stored as a `voice_clones` resource in the user
    repository. Subsequent calls with the same reference audio return the
    cached entry without re-contacting the provider.

    Zero-shot providers (Fish Audio) just keep a pointer to the reference
    audio; persistent-voice_id providers (ElevenLabs, PlayHT) would POST
    the sample once and cache the returned voice_id here.
    """

    @property
    def name(self) -> str:
        return "clone_voice"

    @property
    def description(self) -> str:
        return (
            "Register a voice clone from a reference audio sample. Returns "
            "a `name` the agent can pass to `speak` to synthesize text in "
            "that voice. The same reference audio is detected automatically "
            "(hash-based) and the existing clone is reused — no duplicate "
            "creation, no extra cost. 10-30 s of clean speech is the "
            "recommended sample length. The optional `reference_text` "
            "(transcription of the sample) improves synthesis quality on "
            "zero-shot providers. The user must have the right to clone "
            "the voice in the sample."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "User-friendly name for this voice clone (unique per user). Letters, digits, '-' and '_' are kept; other characters become '_'."},
                "reference_audio_url": {"type": "string", "description": "Reference voice sample (HTTP URL or fs://filestore/<id>/<name>). 10-30 s of clean speech recommended."},
                "reference_text": {"type": "string", "description": "Transcription of the reference audio. Optional; improves quality on zero-shot providers (Fish Audio)."},
                "language": {"type": "string", "description": "BCP-47 language tag of the sample (e.g. 'fr', 'en')."},
            },
            "required": ["name", "reference_audio_url"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no voice-clone service available'}"

        name = (arguments.get("name") or "").strip()
        ref_url_raw = arguments.get("reference_audio_url") or ""
        ref_url = self._rewrite(ref_url_raw, service=svc)
        if not name or not ref_url:
            return ("Error: `name` and `reference_audio_url` are required")

        reference_text = arguments.get("reference_text") or ""
        language = arguments.get("language") or ""

        from core import voice_clone_cache as _cache

        provider = _provider_identity(svc)
        provider_version = _provider_version(svc)

        # Download reference audio so we can hash it.
        try:
            import urllib.request
            req = urllib.request.Request(
                ref_url, headers={"User-Agent": "PawFlow-Agent/1.0"})
            with urllib.request.urlopen(req, timeout=60) as resp:  # nosec B310 - reference audio URLs are user-provided HTTP(S) media inputs.
                ref_bytes = resp.read()
                ref_ct = resp.headers.get("Content-Type",
                                           "application/octet-stream")
        except Exception as e:
            return f"Error downloading reference_audio_url: {e}"

        if not ref_bytes:
            return "Error: reference_audio_url returned empty content"

        ref_hash = _cache.hash_audio(ref_bytes)

        # Already cloned with this exact audio on this provider?
        # Pass `provider_version` so entries from an older API contract
        # are treated as stale and re-cloned against the current version.
        existing = _cache.find_by_hash(
            self._user_id, provider, ref_hash, provider_version)
        if existing:
            _cache.touch(self._user_id, existing.get("name", name))
            return (f"Voice clone already exists: name={existing['name']} "
                    f"provider={provider} (cached, no cost). Use with "
                    f"`speak(voice='{existing['name']}', text=...)`.")

        # Store ref audio in FileStore for future synth calls.
        try:
            from core.file_store import FileStore
            store = FileStore.instance()
            filename = ref_url_raw.rstrip("/").split("/")[-1] or "reference.mp3"
            ref_fid = store.store(
                filename=filename,
                content=ref_bytes,
                content_type=ref_ct,
                conversation_id=self._conversation_id or "_voice_cache",
                user_id=self._user_id,
                ttl=0,
                category="voice_clone_ref",
            )
        except Exception as e:
            return f"Error storing reference audio in FileStore: {e}"

        # Paradigm A providers (ElevenLabs, ...) create a voice_id up-front.
        voice_id = ""
        if hasattr(svc, "ensure_voice_id"):
            try:
                voice_id = _voice_id_value(svc.ensure_voice_id(
                    reference_audio_url=ref_url,
                    reference_text=reference_text,
                    name=name,
                    reference_audio_bytes=ref_bytes,
                ))
            except Exception as e:
                logger.warning(
                    "ensure_voice_id failed on %s: %s — falling back to "
                    "stateless mode", provider, e)
                voice_id = ""

        entry = {
            "name": name,
            "provider": provider,
            "provider_version": provider_version,
            "voice_id": voice_id,
            "ref_audio_hash": ref_hash,
            "ref_audio_fid": ref_fid,
            "ref_audio_filename": filename,
            "ref_audio_content_type": ref_ct,
            "ref_audio_size": len(ref_bytes),
            "reference_text": reference_text,
            "language": language,
        }
        try:
            saved = _cache.save(self._user_id, entry)
        except Exception as e:
            return f"Error persisting voice clone: {e}"

        _msg = (f"Voice clone registered: name={saved['name']} "
                f"provider={provider} paradigm="
                f"{'voice_id' if voice_id else 'zero-shot'}. "
                f"Use `speak(voice='{saved['name']}', text=...)` to "
                f"synthesize speech in this voice.")
        return _msg


class SpeakHandler(_CapabilityHandlerBase):
    """Synthesize text with a registered voice clone, returning an audio URL.

    The agent addresses a voice clone by its user-friendly `voice` name
    (registered earlier via `clone_voice`). The result URL can be passed
    directly to `lipsync` or `speech_to_video`.

    A content-addressed FileStore cache is consulted first:
    (ref_audio_hash, text, language, provider) → cached mp3 avoids
    re-calling the provider for identical re-renders.
    """

    @property
    def name(self) -> str:
        return "speak"

    @property
    def description(self) -> str:
        return (
            "Synthesize speech through the active TTS provider. `voice` may "
            "be either a PawFlow voice registered by `clone_voice` or a "
            "provider-native voice name/id such as Supertonic F1, WaveSpeed "
            "Vivian, or an ElevenLabs voice id. Returns a downloadable audio "
            "URL suitable for `lipsync` / `speech_to_video`. Identical inputs "
            "hit a cache and skip the provider call entirely."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "voice": {"type": "string", "description": "PawFlow voice alias registered via `clone_voice`, or provider-native voice name/id."},
                "text": {"type": "string", "description": "Text to synthesize."},
                "language": {"type": "string", "description": "BCP-47 language tag. Provider may ignore if irrelevant."},
                "service": {"type": "string", "description": "Optional TTS service id override."},
                "audio_service": {"type": "string", "description": "Alias for service; optional TTS service id override."},
                "voice_service": {"type": "string", "description": "Alias for service; optional TTS service id override."},
                "model": {"type": "string", "description": "Provider model override when supported."},
                "destination": {"type": "string"},
                "path": {"type": "string"},
            },
            "required": ["text"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        svc, err = self._get_service(arguments)
        if not svc:
            return f"Error: {err or 'no voice-clone service available'}"

        voice = (arguments.get("voice") or "").strip()
        text = arguments.get("text") or ""
        language = arguments.get("language") or ""
        if not text:
            return "Error: `text` is required"

        from core import voice_clone_cache as _cache
        storage_ttl = self._storage_ttl(arguments)

        provider = _provider_identity(svc)
        entry = _cache.get_by_name(self._user_id, voice) if voice else None
        if entry and entry.get("provider") and entry["provider"] != provider:
            return (f"Error: voice clone {voice!r} was created with provider "
                    f"{entry['provider']}, but the active voice-clone service "
                    f"is {provider}. Switch services to re-use this voice.")

        if not entry:
            # Zero-shot voice-clone providers require a registered PawFlow
            # voice because every synthesis needs reference audio. Providers
            # with native voice ids/names opt in via SUPPORTS_NATIVE_TTS_VOICES,
            # and audio TTS providers expose text_to_speech/speak directly.
            is_clone_provider = hasattr(svc, "clone_speak")
            supports_native = bool(getattr(svc, "SUPPORTS_NATIVE_TTS_VOICES", False))
            has_direct_tts = hasattr(svc, "text_to_speech") or not is_clone_provider
            if voice and is_clone_provider and not supports_native and not has_direct_tts:
                return (f"Error: unknown voice clone {voice!r}. Register it first "
                        f"with `clone_voice`, or select a TTS provider that "
                        f"supports provider-native voices.")
            return self._speak_native_voice(
                svc, provider, voice, text, language, arguments, _cache)

        ref_hash = entry.get("ref_audio_hash") or ""
        cache_key = _cache.tts_cache_key(
            ref_hash, text, language=language, provider=provider)

        # Cache hit — return existing FileStore entry.
        cached_fid = _cache.tts_find(
            self._user_id, self._conversation_id, cache_key,
            include_transient=storage_ttl > 0)
        if cached_fid:
            from core.file_store import FileStore
            meta = FileStore.instance().get_metadata(cached_fid) or {}
            filename = meta.get("filename", f"{voice}.mp3")
            url = f"fs://filestore/{cached_fid}/{filename}"
            _cache.touch(self._user_id, voice)
            return (f"Speech synthesized (cached): {url}\n"
                    f"file_id: {cached_fid}")

        # Cache miss — fetch reference bytes and call the provider.
        ref_fid = entry.get("ref_audio_fid") or ""
        ref_bytes = b""
        if ref_fid:
            try:
                from core.file_store import FileStore
                triple = FileStore.instance().get_required(
                    ref_fid, self._user_id,
                    self._conversation_id or "_voice_cache")
                _, ref_bytes, _ = triple
            except Exception as e:
                logger.warning(
                    "speak: cannot read ref audio %s (%s) — falling back "
                    "to URL fetch", ref_fid, e)
                ref_bytes = b""

        try:
            kwargs = {}
            if entry.get("voice_id"):
                kwargs["voice_id"] = _voice_id_value(entry["voice_id"])
            # Rebuild a usable URL when the service cannot use cached bytes
            # directly (for example WaveSpeedAI prediction inputs expect a
            # URL string, not a raw/base64 sample in this call path).
            ref_url = ""
            if (not ref_bytes or getattr(svc, "REQUIRES_REFERENCE_AUDIO_URL", False)) and ref_fid:
                # Share the reference sample publicly for the duration of
                # this call (revoked when execute() returns).
                ref_url = self._rewrite(f"fs://filestore/{ref_fid}", service=svc)
            r = svc.clone_speak(
                text=text,
                reference_audio_url=ref_url,
                reference_text=entry.get("reference_text") or "",
                language=language,
                reference_audio_bytes=ref_bytes or None,
                **kwargs,
            )
        except Exception as e:
            # Paradigm A: if the provider reports the voice_id is gone
            # (404 / 410 / not_found), our cached entry is stale. Cascade-
            # delete local state and ask the caller to re-register so the
            # next `speak` call doesn't hit the same dead voice_id.
            msg = str(e)
            stale = (
                entry.get("voice_id")
                and (" 404" in msg or " 410" in msg
                     or "not_found" in msg.lower()
                     or "voice_not_found" in msg.lower()))
            if stale:
                try:
                    _cache.cascade_delete(self._user_id, voice, service=None)
                except Exception as de:
                    logger.warning(
                        "speak: cascade_delete after stale voice_id failed: %s",
                        de)
                return (
                    f"Error: voice clone {voice!r} no longer exists on "
                    f"provider {provider} (got: {e}). Local entry purged "
                    f"\u2014 re-register via `clone_voice` before calling "
                    f"`speak` again.")
            return f"Error synthesizing speech: {e}"

        audio_bytes = r.get("audio_bytes") or r.get("bytes") or b""
        audio_path = r.get("audio_path") or r.get("path") or ""
        if not audio_bytes and not audio_path:
            return "Error: provider returned no audio"
        content_type = r.get("content_type", "audio/mpeg")
        ext = {
            "audio/mpeg": "mp3", "audio/mp3": "mp3",
            "audio/wav": "wav", "audio/x-wav": "wav",
            "audio/ogg": "ogg", "audio/L16": "pcm",
        }.get(content_type.split(";")[0].strip(), "mp3")
        filename = arguments.get("path") or (
            f"{voice}_{int(time.time())}.{ext}")

        # Store the rendered audio in FileStore + index it in the TTS cache.
        conv = self._conversation_id or "_voice_cache"
        try:
            if audio_path:
                try:
                    fid = _cache.tts_store_file(
                        user_id=self._user_id,
                        conversation_id=conv,
                        cache_key=cache_key,
                        filename=filename,
                        source_path=str(audio_path),
                        content_type=content_type,
                        ref_audio_hash=ref_hash,
                        ttl=storage_ttl,
                    )
                finally:
                    if r.get("_delete_media_path"):
                        try:
                            os.unlink(str(audio_path))
                        except OSError:
                            pass
            else:
                fid = _cache.tts_store(
                    user_id=self._user_id,
                    conversation_id=conv,
                    cache_key=cache_key,
                    filename=filename,
                    audio_bytes=audio_bytes,
                    content_type=content_type,
                    ref_audio_hash=ref_hash,
                    ttl=storage_ttl,
                )
        except Exception as e:
            return f"Error storing synthesized audio: {e}"

        _cache.touch(self._user_id, voice)
        url = f"fs://filestore/{fid}/{filename}"
        return f"Speech synthesized: {url}\nfile_id: {fid}"

    @staticmethod
    def _storage_ttl(arguments: Dict[str, Any]) -> int:
        try:
            return max(0, int(arguments.get("_tts_storage_ttl") or 0))
        except (TypeError, ValueError):
            return 0

    def _speak_native_voice(self, svc, provider: str, voice: str, text: str,
                            language: str, arguments: Dict[str, Any],
                            cache) -> str:
        """Synthesize with a provider-native voice through BaseTTSService."""
        if not hasattr(svc, "speak"):
            return f"Error: provider {provider} does not support speak()"

        import hashlib
        import json

        kwargs = {k: v for k, v in arguments.items()
                  if k not in ("destination", "path", "text", "voice",
                               "language", *_SERVICE_ARG_NAMES,
                               *_INTERNAL_TTS_ARG_NAMES)}
        storage_ttl = self._storage_ttl(arguments)
        provider_version = _provider_version(svc)
        sig = json.dumps({
            "provider": provider,
            "version": provider_version,
            "voice": voice,
            "params": kwargs,
        }, sort_keys=True, default=str)
        voice_key = hashlib.sha256(sig.encode("utf-8")).hexdigest()
        cache_key = cache.tts_cache_key(
            voice_key, text, language=language, provider=provider)
        cached_fid = cache.tts_find(
            self._user_id, self._conversation_id, cache_key,
            include_transient=storage_ttl > 0)
        if cached_fid:
            from core.file_store import FileStore
            meta = FileStore.instance().get_metadata(cached_fid) or {}
            filename = meta.get("filename", "speech.mp3")
            url = f"fs://filestore/{cached_fid}/{filename}"
            return (f"Speech synthesized (cached): {url}\n"
                    f"file_id: {cached_fid}")

        try:
            r = svc.speak(text=text, voice=voice, language=language, **kwargs)
        except Exception as e:
            return f"Error synthesizing speech: {e}"

        audio_bytes = r.get("audio_bytes") or r.get("bytes") or b""
        audio_path = r.get("audio_path") or r.get("path") or ""
        if not audio_bytes and not audio_path:
            return "Error: provider returned no audio"
        content_type = r.get("content_type", "audio/mpeg")
        ext = {
            "audio/mpeg": "mp3", "audio/mp3": "mp3",
            "audio/wav": "wav", "audio/x-wav": "wav",
            "audio/ogg": "ogg", "audio/L16": "pcm",
            "audio/flac": "flac",
        }.get(content_type.split(";")[0].strip(), "mp3")
        safe_voice = "".join(c if c.isalnum() or c in ("-", "_") else "_"
                              for c in (voice or "speech"))[:40] or "speech"
        filename = arguments.get("path") or f"{safe_voice}_{int(time.time())}.{ext}"
        conv = self._conversation_id or "_voice_cache"
        try:
            if audio_path:
                try:
                    fid = cache.tts_store_file(
                        user_id=self._user_id,
                        conversation_id=conv,
                        cache_key=cache_key,
                        filename=filename,
                        source_path=str(audio_path),
                        content_type=content_type,
                        ref_audio_hash="",
                        ttl=storage_ttl,
                    )
                finally:
                    if r.get("_delete_media_path"):
                        try:
                            os.unlink(str(audio_path))
                        except OSError:
                            pass
            else:
                fid = cache.tts_store(
                    user_id=self._user_id,
                    conversation_id=conv,
                    cache_key=cache_key,
                    filename=filename,
                    audio_bytes=audio_bytes,
                    content_type=content_type,
                    ref_audio_hash="",
                    ttl=storage_ttl,
                )
        except Exception as e:
            return f"Error storing synthesized audio: {e}"
        url = f"fs://filestore/{fid}/{filename}"
        return f"Speech synthesized: {url}\nfile_id: {fid}"


class DeleteVoiceHandler(_CapabilityHandlerBase):
    """Delete a previously registered voice clone — full cascade.

    Removes the user's `voice_clones` entry, the stored reference-audio
    file, every rendered TTS audio that was cached for this voice, and
    (for paradigm-A providers such as ElevenLabs) the voice_id on the
    provider itself so quota is freed upstream.
    """

    @property
    def name(self) -> str:
        return "delete_voice"

    @property
    def description(self) -> str:
        return (
            "Delete a voice clone registered via `clone_voice`. Frees the "
            "provider voice_id (ElevenLabs et al.), removes the stored "
            "reference audio sample and drops every cached synthesis "
            "rendered from this voice. Irreversible."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "voice": {"type": "string", "description": "Name of the voice clone to delete (as returned by `clone_voice`)."},
            },
            "required": ["voice"],
        }

    def execute(self, arguments: Dict[str, Any]) -> str:
        voice = (arguments.get("voice") or "").strip()
        if not voice:
            return "Error: `voice` is required"
        from core import voice_clone_cache as _cache
        # The service may not be configured anymore — we still delete
        # local state, only the provider cleanup step is skipped.
        svc, _ = self._get_service(arguments)
        entry = _cache.get_by_name(self._user_id, voice)
        if entry is None:
            return f"Error: unknown voice clone {voice!r}"
        result = _cache.cascade_delete(self._user_id, voice, svc)
        parts = [f"Voice clone {voice!r} deleted."]
        if result["voice_id"]:
            parts.append("Provider voice_id freed.")
        if result["ref_audio"]:
            parts.append("Reference audio removed.")
        if result["tts_cached"]:
            parts.append(f"{result['tts_cached']} cached rendering(s) purged.")
        return " ".join(parts)
