# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""Image Generation Skill — text-to-image via multiple providers.

Supports:
- Gemini/Nano Banana (Google): gemini-2.5-flash-image, imagen-4.0
- OpenAI DALL-E 3: dall-e-3
- ComfyUI (local): FLUX.2, SDXL, any installed checkpoint
- Automatic1111 (local): Stable Diffusion models

Provider auto-detection: tries configured providers in order.
Generated images auto-save to gallery with optional registration.
"""

from __future__ import annotations

import base64
import json
import logging
import os
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# Gallery output directory
GALLERY_DIR = Path.home() / ".familiar" / "data" / "gallery"
GENERATED_DIR = GALLERY_DIR / "generated"

SKILL_NAME = "image_gen"
SKILL_DESCRIPTION = "Generate images from text prompts using AI"
SKILL_CATEGORY = "creative"

PREREQUISITES = []  # Works with any configured image provider

TOOLS = [
    {
        "name": "generate_image",
        "description": (
            "Generate an image from a text prompt. Supports multiple providers: "
            "Gemini (cloud), DALL-E (cloud), ComfyUI (local), Stable Diffusion (local). "
            "Returns the file path to the generated image."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Text description of the image to generate",
                },
                "style": {
                    "type": "string",
                    "description": "Art style hint (e.g. 'oil painting', 'photograph', 'anime', 'watercolor')",
                    "default": "",
                },
                "provider": {
                    "type": "string",
                    "description": "Provider to use: 'gemini', 'openai', 'comfyui', 'auto' (default: auto-detect)",
                    "default": "auto",
                },
                "width": {
                    "type": "integer",
                    "description": "Image width in pixels (default: 1024)",
                    "default": 1024,
                },
                "height": {
                    "type": "integer",
                    "description": "Image height in pixels (default: 1024)",
                    "default": 1024,
                },
                "register_in_gallery": {
                    "type": "boolean",
                    "description": "Auto-register in Gallery with provenance chain (default: false)",
                    "default": False,
                },
            },
            "required": ["prompt"],
        },
    },
    {
        "name": "list_image_providers",
        "description": "List available image generation providers and their status.",
        "input_schema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "list_comfyui_models",
        "description": (
            "List available models in ComfyUI — checkpoints, LoRAs, samplers. "
            "Only works when ComfyUI is running locally."
        ),
        "input_schema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "list_sd_models",
        "description": (
            "List available Stable Diffusion models in Automatic1111 — "
            "checkpoints, LoRAs, upscalers. Only works when A1111 is running."
        ),
        "input_schema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "upscale_image",
        "description": "Upscale an image using Stable Diffusion's Real-ESRGAN. Doubles resolution by default.",
        "input_schema": {
            "type": "object",
            "properties": {
                "image_path": {
                    "type": "string",
                    "description": "Path to the image to upscale",
                },
                "scale": {
                    "type": "number",
                    "description": "Scale factor (default: 2.0)",
                    "default": 2.0,
                },
                "upscaler": {
                    "type": "string",
                    "description": "Upscaler model (default: R-ESRGAN 4x+)",
                    "default": "R-ESRGAN 4x+",
                },
            },
            "required": ["image_path"],
        },
    },
    {
        "name": "transform_image",
        "description": "Transform an existing image using a text prompt (img2img). Applies AI modifications while preserving structure.",
        "input_schema": {
            "type": "object",
            "properties": {
                "image_path": {
                    "type": "string",
                    "description": "Path to the source image",
                },
                "prompt": {
                    "type": "string",
                    "description": "What to change about the image",
                },
                "strength": {
                    "type": "number",
                    "description": "How much to change (0.0 = no change, 1.0 = complete redraw). Default: 0.75",
                    "default": 0.75,
                },
            },
            "required": ["image_path", "prompt"],
        },
    },
]


def _detect_providers() -> list[dict[str, Any]]:
    """Detect which image generation providers are available."""
    providers = []

    # Gemini
    gemini_key = os.environ.get("GEMINI_API_KEY", "")
    if gemini_key:
        providers.append({
            "name": "gemini",
            "display": "Gemini (Nano Banana)",
            "available": True,
            "models": ["gemini-2.5-flash-image", "gemini-3-pro-image-preview"],
            "note": "Cloud API — requires Gemini API quota",
        })

    # OpenAI DALL-E
    openai_key = os.environ.get("OPENAI_API_KEY", "")
    if openai_key:
        providers.append({
            "name": "openai",
            "display": "OpenAI DALL-E 3",
            "available": True,
            "models": ["dall-e-3", "dall-e-2"],
            "note": "Cloud API — ~$0.04/image (1024x1024)",
        })

    # ComfyUI (local)
    comfyui_url = os.environ.get("COMFYUI_URL", "http://localhost:8188")
    try:
        import httpx
        resp = httpx.get(f"{comfyui_url}/system_stats", timeout=3)
        if resp.status_code == 200:
            providers.append({
                "name": "comfyui",
                "display": "ComfyUI (Local)",
                "available": True,
                "url": comfyui_url,
                "note": "Local — free, requires GPU",
            })
    except Exception:
        pass

    # Automatic1111 (local)
    a1111_url = os.environ.get("A1111_URL", "http://localhost:7860")
    try:
        import httpx
        resp = httpx.get(f"{a1111_url}/sdapi/v1/sd-models", timeout=3)
        if resp.status_code == 200:
            models = [m.get("title", "?") for m in resp.json()[:5]]
            providers.append({
                "name": "automatic1111",
                "display": "Stable Diffusion (A1111)",
                "available": True,
                "url": a1111_url,
                "models": models,
                "note": "Local — free, requires GPU",
            })
    except Exception:
        pass

    return providers


def _save_image(image_bytes: bytes, prompt: str, provider: str, ext: str = "png") -> dict:
    """Save generated image to gallery/generated/ with metadata."""
    GENERATED_DIR.mkdir(parents=True, exist_ok=True)

    image_id = str(uuid.uuid4())[:12]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{image_id}.{ext}"
    filepath = GENERATED_DIR / filename

    filepath.write_bytes(image_bytes)

    # Generate thumbnail
    thumb_path = None
    try:
        from familiar.skills.gallery.watermark import generate_thumbnail

        thumb_dir = GALLERY_DIR / "thumbnails"
        thumb_dir.mkdir(parents=True, exist_ok=True)
        thumb_path = str(thumb_dir / f"{image_id}_thumb.png")
        generate_thumbnail(str(filepath), output_path=thumb_path)
    except Exception as e:
        logger.debug("Thumbnail generation skipped: %s", e)

    # Save metadata
    meta = {
        "image_id": image_id,
        "filename": filename,
        "filepath": str(filepath),
        "thumbnail": thumb_path,
        "prompt": prompt,
        "provider": provider,
        "created_at": datetime.now().isoformat(),
        "size_bytes": len(image_bytes),
    }

    meta_path = GENERATED_DIR / f"{image_id}_meta.json"
    meta_path.write_text(json.dumps(meta, indent=2))

    return meta


def _generate_gemini(prompt: str, style: str, width: int, height: int) -> bytes:
    """Generate image using Gemini API."""
    from google import genai
    from google.genai import types

    api_key = os.environ.get("GEMINI_API_KEY", "")
    client = genai.Client(api_key=api_key)

    full_prompt = f"{prompt}. Style: {style}" if style else prompt

    # Try gemini-2.5-flash-image (native image generation)
    for model in ["gemini-2.5-flash-image", "gemini-3-pro-image-preview"]:
        try:
            response = client.models.generate_content(
                model=model,
                contents=full_prompt,
                config=types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"],
                ),
            )
            for part in response.candidates[0].content.parts:
                if hasattr(part, "inline_data") and part.inline_data:
                    data = part.inline_data.data
                    if isinstance(data, str):
                        return base64.b64decode(data)
                    return data
        except Exception as e:
            logger.debug("Gemini %s failed: %s", model, e)
            continue

    # Fallback: try Imagen
    try:
        response = client.models.generate_images(
            model="imagen-4.0-generate-001",
            prompt=full_prompt,
            config=types.GenerateImagesConfig(number_of_images=1),
        )
        if response.generated_images:
            return response.generated_images[0].image.image_bytes
    except Exception as e:
        logger.debug("Imagen fallback failed: %s", e)

    raise RuntimeError("All Gemini image models failed — check API quota")


def _generate_openai(prompt: str, style: str, width: int, height: int) -> bytes:
    """Generate image using OpenAI DALL-E 3."""
    import httpx

    api_key = os.environ.get("OPENAI_API_KEY", "")
    full_prompt = f"{prompt}. Style: {style}" if style else prompt

    # DALL-E 3 sizes: 1024x1024, 1792x1024, 1024x1792
    size = "1024x1024"
    if width > height * 1.3:
        size = "1792x1024"
    elif height > width * 1.3:
        size = "1024x1792"

    resp = httpx.post(
        "https://api.openai.com/v1/images/generations",
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        json={
            "model": "dall-e-3",
            "prompt": full_prompt,
            "n": 1,
            "size": size,
            "response_format": "b64_json",
        },
        timeout=60,
    )
    resp.raise_for_status()
    data = resp.json()
    b64 = data["data"][0]["b64_json"]
    return base64.b64decode(b64)


def _generate_comfyui(prompt: str, style: str, width: int, height: int) -> bytes:
    """Generate image using local ComfyUI via deep integration module."""
    from familiar.skills.image_gen.comfyui import generate

    full_prompt = f"{prompt}, {style}" if style else prompt
    return generate(prompt=full_prompt, width=width, height=height)


def _generate_comfyui_legacy(prompt: str, style: str, width: int, height: int) -> bytes:
    """Legacy ComfyUI generation — kept as fallback."""
    import httpx

    comfyui_url = os.environ.get("COMFYUI_URL", "http://localhost:8188")
    full_prompt = f"{prompt}, {style}" if style else prompt

    # Simple txt2img workflow
    workflow = {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(time.time()) % 2**32,
                "steps": 20,
                "cfg": 7.5,
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ""}},
        "5": {"class_type": "EmptyLatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}},
        "6": {"class_type": "CLIPTextEncode", "inputs": {"text": full_prompt, "clip": ["4", 1]}},
        "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "ugly, blurry, low quality", "clip": ["4", 1]}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
        "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "familiar", "images": ["8", 0]}},
    }

    # Get first available checkpoint
    try:
        info_resp = httpx.get(f"{comfyui_url}/object_info/CheckpointLoaderSimple", timeout=5)
        ckpts = info_resp.json()["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]
        if ckpts:
            workflow["4"]["inputs"]["ckpt_name"] = ckpts[0]
    except Exception:
        raise RuntimeError("No checkpoints found in ComfyUI")

    # Queue prompt
    resp = httpx.post(f"{comfyui_url}/prompt", json={"prompt": workflow}, timeout=10)
    resp.raise_for_status()
    prompt_id = resp.json()["prompt_id"]

    # Poll for completion
    for _ in range(120):  # 2 min timeout
        time.sleep(1)
        history = httpx.get(f"{comfyui_url}/history/{prompt_id}", timeout=5).json()
        if prompt_id in history:
            outputs = history[prompt_id].get("outputs", {})
            for node_id, output in outputs.items():
                images = output.get("images", [])
                if images:
                    img_info = images[0]
                    img_resp = httpx.get(
                        f"{comfyui_url}/view",
                        params={"filename": img_info["filename"], "subfolder": img_info.get("subfolder", ""), "type": img_info.get("type", "output")},
                        timeout=10,
                    )
                    return img_resp.content
            break

    raise RuntimeError("ComfyUI generation timed out")


# ── Tool handlers ────────────────────────────────────────────────────


def generate_image(data: dict) -> str:
    """Generate an image from a text prompt."""
    prompt = data.get("prompt", "")
    if not prompt:
        return "⚠️ Please provide a prompt describing the image you want."

    style = data.get("style", "")
    provider = data.get("provider", "auto")
    width = data.get("width", 1024)
    height = data.get("height", 1024)
    register = data.get("register_in_gallery", False)

    # Auto-detect provider — check routing config first, then fall back to detection
    if provider == "auto":
        # Check if user set a preferred image provider via routing table
        try:
            import yaml
            config_path = Path.home() / ".familiar" / "config.yaml"
            if config_path.exists():
                with open(config_path) as f:
                    cfg = yaml.safe_load(f) or {}
                # Check routing manual routes for 'image' category
                routes = cfg.get("routing", {}).get("manual_routes", {})
                routed = routes.get("image", "")
                if routed and routed.startswith("image/"):
                    provider = routed.split("/", 1)[1]
                # Legacy: check image_gen.preferred_provider
                elif not routed:
                    pref = cfg.get("image_gen", {}).get("preferred_provider", "")
                    if pref and pref != "auto":
                        provider = pref
        except Exception:
            pass

    if provider == "auto":
        available = _detect_providers()
        if not available:
            return "⚠️ No image generation providers available. Configure Gemini, OpenAI, or install ComfyUI."
        provider = available[0]["name"]

    # Generate
    try:
        if provider == "gemini":
            image_bytes = _generate_gemini(prompt, style, width, height)
        elif provider == "openai":
            image_bytes = _generate_openai(prompt, style, width, height)
        elif provider == "comfyui":
            image_bytes = _generate_comfyui(prompt, style, width, height)
        elif provider in ("automatic1111", "sd", "stable_diffusion"):
            image_bytes = _generate_a1111(prompt, style, width, height)
        else:
            return f"⚠️ Unknown provider: {provider}. Use: gemini, openai, comfyui, sd"
    except Exception as e:
        return f"⚠️ Image generation failed ({provider}): {e}"

    # Save
    meta = _save_image(image_bytes, prompt, provider)

    # Optional gallery registration
    if register:
        try:
            from familiar.skills.gallery.provenance import register_artwork

            result = register_artwork(
                image_path=meta["filepath"],
                title=prompt[:50],
                artist="AI Generated",
                description=f"Prompt: {prompt}\nProvider: {provider}\nStyle: {style or 'default'}",
                tags=["ai-generated", provider, style.replace(" ", "-")] if style else ["ai-generated", provider],
            )
            meta["gallery_id"] = result.get("artwork_id", "")
            meta["registered"] = True
        except Exception as e:
            logger.warning("Gallery registration failed: %s", e)

    return (
        f"✅ Image generated!\n"
        f"  Provider: {provider}\n"
        f"  Size: {meta['size_bytes']:,} bytes\n"
        f"  File: {meta['filepath']}\n"
        f"  Prompt: {prompt[:80]}"
        + (f"\n  Gallery ID: {meta.get('gallery_id', '')}" if meta.get("registered") else "")
    )


def list_image_providers(data: dict) -> str:
    """List available image generation providers."""
    providers = _detect_providers()
    if not providers:
        return (
            "⚠️ No image generation providers configured.\n\n"
            "Options:\n"
            "  • Set GEMINI_API_KEY for Gemini/Nano Banana\n"
            "  • Set OPENAI_API_KEY for DALL-E 3\n"
            "  • Run ComfyUI at localhost:8188 for local FLUX.2/SDXL\n"
            "  • Run Automatic1111 at localhost:7860 for Stable Diffusion"
        )

    lines = ["🎨 Available Image Generation Providers:\n"]
    for p in providers:
        models = ", ".join(p.get("models", [])[:3]) if p.get("models") else ""
        lines.append(f"  ✅ {p['display']}")
        if models:
            lines.append(f"     Models: {models}")
        lines.append(f"     {p['note']}")
        lines.append("")

    return "\n".join(lines)


def list_comfyui_models(data: dict) -> str:
    """List available models in a local ComfyUI instance."""
    try:
        from familiar.skills.image_gen.comfyui import (
            is_available, list_checkpoints, list_loras, list_samplers,
            list_schedulers, get_status_summary,
        )
    except ImportError:
        return "⚠️ ComfyUI module not available."

    if not is_available():
        return (
            "⚠️ ComfyUI not running.\n\n"
            "Start ComfyUI at http://localhost:8188 or set COMFYUI_URL.\n"
            "Install: /connect comfyui"
        )

    lines = [get_status_summary(), ""]

    ckpts = list_checkpoints()
    if ckpts:
        lines.append(f"📦 Checkpoints ({len(ckpts)}):")
        for c in ckpts[:10]:
            lines.append(f"  • {c}")
        if len(ckpts) > 10:
            lines.append(f"  ... and {len(ckpts) - 10} more")
        lines.append("")

    loras = list_loras()
    if loras:
        lines.append(f"🎭 LoRAs ({len(loras)}):")
        for l in loras[:10]:
            lines.append(f"  • {l}")
        if len(loras) > 10:
            lines.append(f"  ... and {len(loras) - 10} more")
        lines.append("")

    samplers = list_samplers()
    lines.append(f"⚙️ Samplers: {', '.join(samplers[:8])}")

    schedulers = list_schedulers()
    lines.append(f"📅 Schedulers: {', '.join(schedulers[:6])}")

    return "\n".join(lines)


def _generate_a1111(prompt: str, style: str, width: int, height: int) -> bytes:
    """Generate image using Automatic1111 Stable Diffusion API."""
    from familiar.skills.image_gen.automatic1111 import txt2img

    full_prompt = f"{prompt}, {style}" if style else prompt
    images = txt2img(prompt=full_prompt, width=width, height=height)
    if not images:
        raise RuntimeError("A1111 returned no images")
    return images[0]


def list_sd_models(data: dict) -> str:
    """List available Stable Diffusion models in A1111."""
    try:
        from familiar.skills.image_gen.automatic1111 import (
            is_available, get_status_summary, list_models, list_loras, list_upscalers,
        )
    except ImportError:
        return "⚠️ A1111 module not available."

    if not is_available():
        return (
            "⚠️ Stable Diffusion (A1111) not running.\n\n"
            "Start at http://localhost:7860 or set A1111_URL.\n"
            "Install: https://github.com/AUTOMATIC1111/stable-diffusion-webui"
        )

    lines = [get_status_summary(), ""]

    models = list_models()
    if models:
        lines.append(f"📦 Models ({len(models)}):")
        for m in models[:10]:
            lines.append(f"  • {m['title']}")
        if len(models) > 10:
            lines.append(f"  ... and {len(models) - 10} more")
        lines.append("")

    loras = list_loras()
    if loras:
        lines.append(f"🎭 LoRAs ({len(loras)}):")
        for l in loras[:10]:
            lines.append(f"  • {l['name']}")
        lines.append("")

    upscalers = list_upscalers()
    if upscalers:
        lines.append(f"🔍 Upscalers: {', '.join(upscalers[:6])}")

    return "\n".join(lines)


def upscale_image(data: dict) -> str:
    """Upscale an image using Stable Diffusion's upscaler."""
    image_path = data.get("image_path", "")
    if not image_path:
        return "⚠️ Please provide an image_path."

    path = Path(image_path).expanduser()
    if not path.exists():
        return f"⚠️ File not found: {image_path}"

    scale = data.get("scale", 2.0)
    upscaler = data.get("upscaler", "R-ESRGAN 4x+")

    try:
        from familiar.skills.image_gen.automatic1111 import upscale, is_available

        if not is_available():
            return "⚠️ A1111 not running. Start Stable Diffusion WebUI first."

        image_bytes = path.read_bytes()
        result = upscale(image_bytes, upscaler=upscaler, scale=scale)

        # Save upscaled image
        out_name = f"{path.stem}_upscaled_{scale}x{path.suffix}"
        out_path = path.parent / out_name
        out_path.write_bytes(result)

        return (
            f"✅ Image upscaled {scale}x!\n"
            f"  Original: {path.name} ({len(image_bytes):,} bytes)\n"
            f"  Upscaled: {out_name} ({len(result):,} bytes)\n"
            f"  Saved: {out_path}"
        )
    except Exception as e:
        return f"⚠️ Upscale failed: {e}"


def transform_image(data: dict) -> str:
    """Transform an image using img2img."""
    image_path = data.get("image_path", "")
    prompt = data.get("prompt", "")
    if not image_path or not prompt:
        return "⚠️ Please provide image_path and prompt."

    path = Path(image_path).expanduser()
    if not path.exists():
        return f"⚠️ File not found: {image_path}"

    strength = data.get("strength", 0.75)

    try:
        from familiar.skills.image_gen.automatic1111 import img2img, is_available

        if not is_available():
            return "⚠️ A1111 not running. Start Stable Diffusion WebUI first."

        image_bytes = path.read_bytes()
        results = img2img(image_bytes, prompt=prompt, denoising_strength=strength)
        if not results:
            return "⚠️ No images returned."

        # Save transformed image
        meta = _save_image(results[0], f"img2img: {prompt}", "automatic1111")
        return (
            f"✅ Image transformed!\n"
            f"  Prompt: {prompt[:60]}\n"
            f"  Strength: {strength}\n"
            f"  File: {meta['filepath']}"
        )
    except Exception as e:
        return f"⚠️ Transform failed: {e}"


# ── Wire handlers into TOOLS dicts (skill loader requires 'handler' key) ──
_HANDLER_MAP = {
    "generate_image": generate_image,
    "list_image_providers": list_image_providers,
    "list_comfyui_models": list_comfyui_models,
    "list_sd_models": list_sd_models,
    "upscale_image": upscale_image,
    "transform_image": transform_image,
}
for _tool in TOOLS:
    if _tool["name"] in _HANDLER_MAP:
        _tool["handler"] = _HANDLER_MAP[_tool["name"]]
