Generation

def edit_image(
    prompt: str,
    image_path: str,
    provider: str = "openai",
    model: str = None,
    height: int = 1024,
    width: int = 1024,
    save_path: Optional[str] = None,
    api_key: Optional[str] = None,
):
    """
    Convenience function to edit an existing image.

    Args:
        prompt (str): The editing instructions.
        image_path (str): Path to the image to edit.
        provider (str): The provider to use for editing.
        model (str): The model to use.
        height (int): The height of the output image.
        width (int): The width of the output image.
        save_path (str): Path to save the edited image.
        api_key (str): API key for the provider.

    Returns:
        PIL.Image: The edited image.
    """
    if model is None:
        if provider == "openai":
            model = "gpt-image-1"
        elif provider == "gemini":
            model = "gemini-2.5-flash-image"

    image = generate_image(
        prompt=prompt,
        provider=provider,
        model=model,
        attachments=[image_path],
        height=height,
        width=width,
        save_path=save_path,
        api_key=api_key
    )

    return image

def gemini_image_gen(
    prompt: str,
    model: str = "gemini-2.5-flash",
    attachments: Union[List[Union[str, bytes, Image.Image]], None] = None,
    n_images: int = 1,
    api_key: Optional[str] = None,
):
    """Generate or edit images using Google's Gemini API."""
    from google import genai
    from google.genai import types
    from io import BytesIO
    import requests

    if api_key is None:
        api_key = os.environ.get('GEMINI_API_KEY')

    client = genai.Client(api_key=api_key)
    collected_images = []

    if attachments is not None:
        processed_contents = [prompt]

        for attachment in attachments:
            if isinstance(attachment, str):
                if attachment.startswith(('http://', 'https://')):
                    image_bytes = requests.get(attachment).content
                    mime_type = 'image/jpeg'
                    processed_contents.append(
                        types.Part.from_bytes(
                            data=image_bytes,
                            mime_type=mime_type
                        )
                    )
                else:
                    with open(attachment, 'rb') as f:
                        image_bytes = f.read()

                    if attachment.lower().endswith('.png'):
                        mime_type = 'image/png'
                    elif attachment.lower().endswith('.jpg') or attachment.lower().endswith('.jpeg'):
                        mime_type = 'image/jpeg'
                    elif attachment.lower().endswith('.webp'):
                        mime_type = 'image/webp'
                    else:
                        mime_type = 'image/jpeg'

                    processed_contents.append(
                        types.Part.from_bytes(
                            data=image_bytes,
                            mime_type=mime_type
                        )
                    )
            elif isinstance(attachment, bytes):
                processed_contents.append(
                    types.Part.from_bytes(
                        data=attachment,
                        mime_type='image/jpeg'
                    )
                )
            elif isinstance(attachment, Image.Image):
                img_byte_arr = BytesIO()
                attachment.save(img_byte_arr, format='PNG')
                img_byte_arr.seek(0)
                processed_contents.append(
                    types.Part.from_bytes(
                        data=img_byte_arr.getvalue(),
                        mime_type='image/png'
                    )
                )

        response = client.models.generate_content(
            model=model,
            contents=processed_contents,
            config=types.GenerateContentConfig(
                response_modalities=["IMAGE", "TEXT"],
            ),
        )

        if hasattr(response, 'candidates') and response.candidates:
            for candidate in response.candidates:
                if candidate.content and candidate.content.parts:
                    for part in candidate.content.parts:
                        if hasattr(part, 'inline_data') and part.inline_data:
                            image_data = part.inline_data.data
                            collected_images.append(Image.open(BytesIO(image_data)))

        if not collected_images and hasattr(response, 'text'):
            print(f"Gemini response text: {response.text}")

        return collected_images
    else:
        if 'imagen' in model:
            response = client.models.generate_images(
                model=model,
                prompt=prompt,
                config=types.GenerateImagesConfig(
                    number_of_images=n_images,
                )
            )
            for generated_image in response.generated_images:
                collected_images.append(Image.open(BytesIO(generated_image.image.image_bytes)))
            return collected_images

        elif 'flash-image' in model or 'image-preview' in model or '2.5-flash' in model:
            response = client.models.generate_content(
                model=model,
                contents=[prompt],
                config=types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"],
                ),
            )

            if hasattr(response, 'candidates') and response.candidates:
                for candidate in response.candidates:
                    if candidate.content and candidate.content.parts:
                        for part in candidate.content.parts:
                            if hasattr(part, 'inline_data') and part.inline_data:
                                image_data = part.inline_data.data
                                collected_images.append(Image.open(BytesIO(image_data)))

            if not collected_images and hasattr(response, 'text'):
                print(f"Gemini response text: {response.text}")

            return collected_images

        else:
            raise ValueError(f"Unsupported Gemini image model or API usage for new generation: '{model}'")

def generate_image(
    prompt: str,
    model: str ,
    provider: str ,
    height: int = 1024,
    width: int = 1024,
    n_images: int = 1,
    api_key: Optional[str] = None,
    api_url: Optional[str] = None,
    attachments: Union[List[Union[str, bytes, Image.Image]], None] = None,
    save_path: Optional[str] = None,
    custom_model_path: Optional[str] = None,

):
    """
    Unified function to generate or edit images using various providers.

    Args:
        prompt (str): The prompt for generating/editing the image.
        model (str): The model to use.
        provider (str): The provider to use ('openai', 'diffusers', 'gemini', 'ollama').
        height (int): The height of the output image.
        width (int): The width of the output image.
        n_images (int): Number of images to generate.
        api_key (str): API key for the provider.
        api_url (str): API URL for the provider.
        attachments (list): List of images for editing. Can be file paths, bytes, or PIL Images.
        save_path (str): Path to save the generated image.
        custom_model_path (str): Path to a locally fine-tuned Diffusers model.

    Returns:
        List[PIL.Image.Image]: A list of generated PIL Image objects.
    """
    from urllib.request import urlopen
    import os

    if model is None and custom_model_path is None:
        if provider == "openai":
            model = "dall-e-2"
        elif provider == "diffusers":
            model = "runwayml/stable-diffusion-v1-5"
        elif provider == "gemini":
            model = "gemini-2.5-flash-image"
        elif provider == "ollama":
            model = "x/z-image-turbo"

    all_generated_pil_images = []

    if provider == "diffusers":
        if custom_model_path and os.path.isdir(custom_model_path):
            print(f"🌋 Using custom Diffusers model from path: {custom_model_path}")
            model_to_use = custom_model_path
        else:
            model_to_use = model
            print(f"🌋 Using standard Diffusers model: {model_to_use}")

        for _ in range(n_images):
            try:
                image = generate_image_diffusers(
                    prompt=prompt, 
                    model=model_to_use,
                    height=height, 
                    width=width
                )
                all_generated_pil_images.append(image)
            except MemoryError as e:
                raise e
            except Exception as e:
                raise e

    elif provider == "openai":
        images = openai_image_gen(
            prompt=prompt,
            model=model,
            attachments=attachments,
            height=height,
            width=width,
            n_images=n_images,
            api_key=api_key

        )
        all_generated_pil_images.extend(images)

    elif provider == "gemini":
        images = gemini_image_gen(
            prompt=prompt,
            model=model,
            attachments=attachments,
            n_images=n_images,
            api_key=api_key
        )
        all_generated_pil_images.extend(images)

    elif provider == "ollama":
        images = ollama_image_gen(
            prompt=prompt,
            model=model,
            height=height,
            width=width,
            n_images=n_images,
            api_url=api_url
        )
        all_generated_pil_images.extend(images)

    else:
        valid_sizes = ["256x256", "512x512", "1024x1024", "1024x1792", "1792x1024"]
        size = f"{width}x{height}"

        if attachments is not None:
            raise ValueError("Image editing not supported with litellm provider")

        litellm_model_string = f"{provider}/{model}" if provider and model else model

        image_response = image_generation(
            prompt=prompt,
            model=litellm_model_string,
            n=n_images,
            size=size,
            api_key=api_key,
            api_base=api_url,
        )
        for item_data in image_response.data:
            if hasattr(item_data, 'url') and item_data.url:
                with urlopen(item_data.url) as response:
                    all_generated_pil_images.append(Image.open(response))
            elif hasattr(item_data, 'b64_json') and item_data.b64_json:
                image_bytes = base64.b64decode(item_data.b64_json)
                all_generated_pil_images.append(Image.open(io.BytesIO(image_bytes)))
            else:
                print(f"Warning: litellm ImageResponse item has no URL or b64_json: {item_data}")

    if save_path:
        for i, img_item in enumerate(all_generated_pil_images):
            temp_save_path = f"{os.path.splitext(save_path)[0]}_{i}{os.path.splitext(save_path)[1]}"
            if isinstance(img_item, Image.Image):
                img_item.save(temp_save_path)
            else:
                print(f"Warning: Attempting to save non-PIL image item: {type(img_item)}. Skipping save for this item.")

    return all_generated_pil_images

def generate_image_diffusers(
    prompt: str,
    model: str = "runwayml/stable-diffusion-v1-5",
    device: str = "cpu",
    height: int = 512,
    width: int = 512,
):
    """Generate an image using the Stable Diffusion API with memory optimization."""
    import torch
    import gc
    import os
    from diffusers import DiffusionPipeline, StableDiffusionPipeline

    try:
        torch_dtype = torch.float16 if device != "cpu" and torch.cuda.is_available() else torch.float32

        if os.path.isdir(model):
            print(f"🌋 Loading fine-tuned Diffusers model from local path: {model}")

            checkpoint_path = os.path.join(model, 'model_final.pt')
            if os.path.exists(checkpoint_path):
                print(f"🌋 Found model_final.pt at {checkpoint_path}.")

                checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)

                cfg = checkpoint.get('config')
                has_image_size = (isinstance(cfg, dict) and 'image_size' in cfg) or hasattr(cfg, 'image_size')
                if 'config' in checkpoint and has_image_size:
                    print(f"🌋 Detected custom SimpleUNet model, using custom generation")
                    from npcpy.ft.diff import generate_image as custom_generate_image

                    image = custom_generate_image(
                        model_path=checkpoint_path,
                        prompt=prompt,
                        num_samples=1,
                        image_size=height
                    )
                    return image

                else:
                    print(f"🌋 Detected Stable Diffusion UNet checkpoint")
                    base_model_id = "runwayml/stable-diffusion-v1-5"
                    print(f"🌋 Loading base pipeline: {base_model_id}")
                    pipe = StableDiffusionPipeline.from_pretrained(
                        base_model_id,
                        torch_dtype=torch_dtype,
                        use_safetensors=True,
                        variant="fp16" if torch_dtype == torch.float16 else None,
                    )

                    print(f"🌋 Loading custom UNet weights from {checkpoint_path}")

                    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
                        unet_state_dict = checkpoint['model_state_dict']
                        print(f"🌋 Extracted model_state_dict from checkpoint")
                    else:
                        unet_state_dict = checkpoint
                        print(f"🌋 Using checkpoint directly as state_dict")

                    pipe.unet.load_state_dict(unet_state_dict)
                    pipe = pipe.to(device)
                    print(f"🌋 Successfully loaded fine-tuned UNet weights")

            else:
                raise OSError(f"Error: Fine-tuned model directory {model} does not contain 'model_final.pt'")

        else:
            print(f"🌋 Loading standard Diffusers model: {model}")
            if 'Qwen' in model:
                pipe = DiffusionPipeline.from_pretrained(
                    model,
                    torch_dtype=torch_dtype,
                    use_safetensors=True,
                    variant="fp16" if torch_dtype == torch.float16 else None,
                )
            else:
                pipe = DiffusionPipeline.from_pretrained(
                    model,
                    torch_dtype=torch_dtype,
                    use_safetensors=True,
                    variant="fp16" if torch_dtype == torch.float16 else None,
                )

        if hasattr(pipe, 'enable_attention_slicing'):
            pipe.enable_attention_slicing()

        if hasattr(pipe, 'enable_model_cpu_offload') and device != "cpu":
            pipe.enable_model_cpu_offload()
        else:
            pipe = pipe.to(device)

        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
            try:
                pipe.enable_xformers_memory_efficient_attention()
            except Exception:
                pass

        with torch.inference_mode():
            result = pipe(prompt, height=height, width=width, num_inference_steps=20)
            image = result.images[0]

        del pipe
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

        return image

    except Exception as e:
        if 'pipe' in locals():
            del pipe
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

        if "out of memory" in str(e).lower() or "killed" in str(e).lower():
            print(f"Memory error during image generation: {e}")
            print("Suggestions:")
            print("1. Try using a smaller model like 'CompVis/stable-diffusion-v1-4'")
            print("2. Reduce image dimensions (e.g., height=256, width=256)")
            print("3. Use a cloud service for image generation")
            raise MemoryError(f"Insufficient memory for image generation with model {model}. Try a smaller model or reduce image size.")
        else:
            raise e

def ollama_image_gen(
    prompt: str,
    model: str = "x/z-image-turbo",
    height: int = 512,
    width: int = 512,
    n_images: int = 1,
    api_url: Optional[str] = None,
    seed: Optional[int] = None,
    negative_prompt: Optional[str] = None,
    num_steps: Optional[int] = None,
):
    """Generate images using Ollama's image generation API.

    Works with ollama image gen models like x/z-image-turbo and x/flux2-klein.
    Uses the /api/generate endpoint with image gen specific options.
    """
    import requests

    if api_url is None:
        api_url = os.environ.get('OLLAMA_API_URL', 'http://localhost:11434')

    endpoint = f"{api_url}/api/generate"

    collected_images = []

    for _ in range(n_images):
        options = {}
        if width:
            options["width"] = width
        if height:
            options["height"] = height
        if seed is not None:
            options["seed"] = seed
        if num_steps is not None:
            options["num_steps"] = num_steps

        payload = {
            "model": model,
            "prompt": prompt,
            "stream": False,
        }
        if options:
            payload["options"] = options
        if negative_prompt:
            payload["negative_prompt"] = negative_prompt

        response = requests.post(endpoint, json=payload)

        if not response.ok:
            try:
                err = response.json()
                err_msg = err.get('error', response.text)
            except Exception:
                err_msg = response.text
            raise RuntimeError(f"Ollama image gen failed ({response.status_code}): {err_msg}\nModel: {model} — make sure it's pulled (`ollama pull {model}`)")

        result = response.json()

        if 'image' in result and result['image']:
            image_bytes = base64.b64decode(result['image'])
            image = Image.open(io.BytesIO(image_bytes))
            collected_images.append(image)
        elif 'images' in result and result['images']:
            for img_b64 in result['images']:
                image_bytes = base64.b64decode(img_b64)
                image = Image.open(io.BytesIO(image_bytes))
                collected_images.append(image)
        else:
            raise ValueError(f"No images returned from Ollama. Response keys: {list(result.keys())}. Make sure '{model}' is an image generation model (e.g. x/z-image-turbo, x/flux2-klein).")

    return collected_images

def openai_image_gen(
    prompt: str,
    model: str = "dall-e-2",
    attachments: Union[List[Union[str, bytes, Image.Image]], None] = None,
    height: int = 1024,
    width: int = 1024,
    n_images: int = 1,
    api_key: Optional[str] = None,
):
    """Generate or edit an image using the OpenAI API."""
    from openai import OpenAI

    client = OpenAI(api_key=api_key) if api_key else OpenAI()

    if height is None:
        height = 1024
    if width is None:
        width = 1024

    size_str = f"{width}x{height}"

    if attachments is not None:
        processed_images = []
        files_to_close = []
        for attachment in attachments:
            if isinstance(attachment, str):
                file_handle = open(attachment, "rb")
                processed_images.append(file_handle)
                files_to_close.append(file_handle)
            elif isinstance(attachment, bytes):
                img_byte_arr = io.BytesIO(attachment)
                img_byte_arr.name = 'image.png'
                processed_images.append(img_byte_arr)
            elif isinstance(attachment, Image.Image):
                img_byte_arr = io.BytesIO()
                attachment.save(img_byte_arr, format='PNG')
                img_byte_arr.seek(0)
                img_byte_arr.name = 'image.png'
                processed_images.append(img_byte_arr)

        try:
            result = client.images.edit(
                model=model,
                image=processed_images[0],
                prompt=prompt,
                n=n_images,
                size=size_str,
            )
        finally:
            for f in files_to_close:
                f.close()
    else:
        result = client.images.generate(
            model=model,
            prompt=prompt,
            n=n_images,
            size=size_str,
        )

    collected_images = []
    for item_data in result.data:
        b64 = getattr(item_data, 'b64_json', None)
        url = getattr(item_data, 'url', None)
        if b64:
            image_bytes = base64.b64decode(b64)
            image = Image.open(io.BytesIO(image_bytes))
        elif url:
            response = requests.get(url)
            response.raise_for_status()
            image = Image.open(io.BytesIO(response.content))
        else:
            raise ValueError(
                f"OpenAI image response had neither b64_json nor url for model={model}: {item_data!r}"
            )
        collected_images.append(image)

    return collected_images

def _get_qwen3_model(
    model_size: str = "1.7B",
    model_type: str = "custom_voice",
    device: str = "auto",
):
    """Load and cache a Qwen3-TTS model."""
    cache_key = (model_size, model_type, device)
    if cache_key in _qwen3_model_cache:
        return _qwen3_model_cache[cache_key]

    import torch
    from huggingface_hub import snapshot_download

    if device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

    dtype = torch.bfloat16 if device != "cpu" else torch.float32

    size_tag = "0.6B" if "0.6" in model_size else "1.7B"
    type_map = {
        "custom_voice": f"Qwen/Qwen3-TTS-12Hz-{size_tag}-CustomVoice",
        "base": f"Qwen/Qwen3-TTS-12Hz-{size_tag}-Base",
        "voice_design": f"Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
    }

    repo_id = type_map.get(model_type, type_map["custom_voice"])

    cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "qwen-tts")
    model_dir = os.path.join(cache_dir, repo_id.split("/")[-1])

    if not os.path.exists(os.path.join(model_dir, "config.json")):
        os.makedirs(cache_dir, exist_ok=True)
        snapshot_download(repo_id=repo_id, local_dir=model_dir)

    try:
        from qwen_tts import Qwen3TTSModel
    except ImportError:
        raise ImportError("qwen_tts package not found. Install from: https://github.com/QwenLM/Qwen3-TTS or pip install qwen-tts")

    model = Qwen3TTSModel.from_pretrained(
        model_dir, device_map=device, dtype=dtype
    )

    _qwen3_model_cache.clear()
    _qwen3_model_cache[cache_key] = model
    return model

def _music_one(provider: str, prompt: str, model: Optional[str], duration: int, api_key: Optional[str]) -> dict:
    p = (provider or "").lower()
    if p in ("local", "musicgen", "transformers", "meta"):
        m = model or "facebook/musicgen-small"
        return {"audio": music_musicgen_local(prompt, duration=duration, model=m),
                "format": "wav", "provider": "musicgen-local", "model": m}
    if p in ("replicate",):
        m = model or "meta/musicgen"
        return {"audio": music_replicate(prompt, duration=duration, model=m, api_key=api_key),
                "format": "wav", "provider": "replicate", "model": m}
    if p in ("elevenlabs", "stability", "11labs"):
        return {"audio": music_elevenlabs_sfx(prompt, duration=duration, api_key=api_key),
                "format": "mp3", "provider": "elevenlabs", "model": "sound-generation"}
    raise ValueError(f"Unknown music provider: {provider!r}")

def audio_to_base64(audio_data: bytes) -> str:
    """Encode audio to base64 string."""
    return base64.b64encode(audio_data).decode('utf-8')

def base64_to_audio(b64_string: str) -> bytes:
    """Decode base64 string to audio bytes."""
    return base64.b64decode(b64_string)

async def gemini_live_connect(
    api_key: Optional[str] = None,
    model: str = "gemini-2.0-flash-exp",
    voice: str = "Puck",
    system_instruction: str = "You are a helpful assistant."
):
    """
    Connect to Gemini Live API.

    Returns:
        WebSocket connection
    """
    import websockets

    api_key = api_key or os.environ.get('GOOGLE_API_KEY') or os.environ.get('GEMINI_API_KEY')

    url = f"wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContent?key={api_key}"

    ws = await websockets.connect(url)

    await ws.send(json.dumps({
        "setup": {
            "model": f"models/{model}",
            "generation_config": {
                "response_modalities": ["AUDIO"],
                "speech_config": {
                    "voice_config": {
                        "prebuilt_voice_config": {"voice_name": voice}
                    }
                }
            },
            "system_instruction": {"parts": [{"text": system_instruction}]}
        }
    }))

    response = await ws.recv()
    data = json.loads(response)
    if "setupComplete" not in data:
        await ws.close()
        raise Exception(f"Gemini Live setup failed: {data}")

    return ws

async def gemini_live_receive(ws, on_audio=None, on_text=None):
    """
    Receive response from Gemini Live.

    Returns:
        Tuple of (full_audio_bytes, full_text)
    """
    audio_chunks = []
    text_chunks = []

    async for message in ws:
        data = json.loads(message)

        if "serverContent" in data:
            content = data["serverContent"]

            if "modelTurn" in content:
                for part in content["modelTurn"].get("parts", []):
                    if "inlineData" in part:
                        audio = base64.b64decode(part["inlineData"].get("data", ""))
                        audio_chunks.append(audio)
                        if on_audio:
                            on_audio(audio)
                    elif "text" in part:
                        text_chunks.append(part["text"])
                        if on_text:
                            on_text(part["text"])

            if content.get("turnComplete"):
                break

    return b''.join(audio_chunks), ''.join(text_chunks)

async def gemini_live_send_audio(ws, audio_data: bytes, mime_type: str = "audio/pcm"):
    """Send audio to Gemini Live."""
    await ws.send(json.dumps({
        "realtime_input": {
            "media_chunks": [{
                "data": base64.b64encode(audio_data).decode(),
                "mime_type": mime_type
            }]
        }
    }))

async def gemini_live_send_text(ws, text: str):
    """Send text message to Gemini Live."""
    await ws.send(json.dumps({
        "client_content": {
            "turns": [{"role": "user", "parts": [{"text": text}]}],
            "turn_complete": True
        }
    }))

def generate_music(
    prompt: str,
    provider: str = "replicate",
    model: Optional[str] = None,
    duration: int = 10,
    api_key: Optional[str] = None,
) -> dict:
    """Unified music-generation entry point with automatic provider fallback.

    Tries the requested provider first; on failure, walks the other providers
    that have credentials available so the user always gets audio back.
    """
    requested = (provider or "local").lower()
    order = [requested]
    for cand in ("local", "replicate", "elevenlabs"):
        if cand not in order:
            order.append(cand)

    errors: list[str] = []
    for prov in order:
        # Skip cloud providers that have no credentials configured.
        if prov == "replicate" and not (api_key or os.environ.get("REPLICATE_API_TOKEN") or os.environ.get("REPLICATE_API_KEY")):
            errors.append(f"{prov}: no REPLICATE_API_TOKEN")
            continue
        if prov in ("elevenlabs", "11labs", "stability") and not (api_key or os.environ.get("ELEVENLABS_API_KEY")):
            errors.append(f"{prov}: no ELEVENLABS_API_KEY")
            continue
        try:
            return _music_one(prov, prompt, model if prov == requested else None, duration, api_key)
        except Exception as e:
            errors.append(f"{prov}: {e}")
            continue

    raise RuntimeError("All music providers failed:\n  - " + "\n  - ".join(errors))

def get_available_engines() -> dict:
    """Get info about available TTS engines."""
    engines = {
        "kokoro": {
            "name": "Kokoro",
            "type": "local",
            "available": False,
            "description": "Local neural TTS (82M params)",
            "install": "pip install kokoro soundfile"
        },
        "qwen3": {
            "name": "Qwen3-TTS",
            "type": "local",
            "available": False,
            "description": "Local high-quality multilingual TTS (0.6B/1.7B)",
            "install": "pip install qwen-tts torch torchaudio transformers"
        },
        "elevenlabs": {
            "name": "ElevenLabs",
            "type": "cloud",
            "available": False,
            "description": "High-quality cloud TTS",
            "requires": "ELEVENLABS_API_KEY"
        },
        "openai": {
            "name": "OpenAI Realtime",
            "type": "cloud",
            "available": False,
            "description": "OpenAI real-time voice API",
            "requires": "OPENAI_API_KEY"
        },
        "gemini": {
            "name": "Gemini Live",
            "type": "cloud",
            "available": False,
            "description": "Google Gemini real-time voice",
            "requires": "GOOGLE_API_KEY or GEMINI_API_KEY"
        },
        "gtts": {
            "name": "Google TTS",
            "type": "cloud",
            "available": False,
            "description": "Free Google TTS (fallback)"
        }
    }

    try:
        from kokoro import KPipeline
        engines["kokoro"]["available"] = True
    except ImportError:
        pass

    try:
        from qwen_tts import Qwen3TTSModel
        engines["qwen3"]["available"] = True
    except ImportError:
        pass

    if os.environ.get('ELEVENLABS_API_KEY'):
        engines["elevenlabs"]["available"] = True

    if os.environ.get('OPENAI_API_KEY'):
        engines["openai"]["available"] = True

    if os.environ.get('GOOGLE_API_KEY') or os.environ.get('GEMINI_API_KEY'):
        engines["gemini"]["available"] = True

    try:
        from gtts import gTTS
        engines["gtts"]["available"] = True
    except ImportError:
        pass

    return engines

def get_available_voices(engine: str = "kokoro") -> list:
    """Get available voices for an engine."""
    engine = engine.lower()

    if engine == "kokoro":
        return get_kokoro_voices()
    elif engine in ("qwen3", "qwen3-tts", "qwen"):
        return get_qwen3_voices()
    elif engine == "elevenlabs":
        return get_elevenlabs_voices()
    elif engine == "openai":
        return get_openai_voices()
    elif engine == "gemini":
        return get_gemini_voices()
    elif engine == "gtts":
        return get_gtts_voices()
    else:
        return []

def get_elevenlabs_voices(api_key: Optional[str] = None) -> list:
    """Get available ElevenLabs voices."""
    if api_key is None:
        api_key = os.environ.get('ELEVENLABS_API_KEY')

    if not api_key:
        return []

    try:
        from elevenlabs.client import ElevenLabs
        client = ElevenLabs(api_key=api_key)
        voices = client.voices.get_all()
        return [{"id": v.voice_id, "name": v.name} for v in voices.voices]
    except Exception:
        return []

def get_gemini_voices() -> list:
    """Get available Gemini Live voices."""
    return [
        {"id": "Puck", "name": "Puck"},
        {"id": "Charon", "name": "Charon"},
        {"id": "Kore", "name": "Kore"},
        {"id": "Fenrir", "name": "Fenrir"},
        {"id": "Aoede", "name": "Aoede"},
    ]