# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# Licensed under the MIT License

"""Automatic1111 (Stable Diffusion WebUI) integration.

Connects to a locally-running A1111 instance for Stable Diffusion
models, LoRAs, upscaling, img2img, and inpainting.
"""

from __future__ import annotations

import base64
import logging
import os
import time
from typing import Any, Optional

logger = logging.getLogger(__name__)

DEFAULT_URL = "http://localhost:7860"


def get_url() -> str:
    return os.environ.get("A1111_URL", DEFAULT_URL)


def is_available() -> bool:
    """Check if Automatic1111 is reachable."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/sdapi/v1/sd-models", timeout=3)
        return resp.status_code == 200
    except Exception:
        return False


def list_models() -> list[dict[str, str]]:
    """List available SD models (checkpoints)."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/sdapi/v1/sd-models", timeout=5)
        if resp.status_code == 200:
            return [{"title": m.get("title", ""), "name": m.get("model_name", "")} for m in resp.json()]
    except Exception:
        pass
    return []


def list_loras() -> list[dict[str, str]]:
    """List available LoRA models."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/sdapi/v1/loras", timeout=5)
        if resp.status_code == 200:
            return [{"name": l.get("name", ""), "alias": l.get("alias", "")} for l in resp.json()]
    except Exception:
        pass
    return []


def list_samplers() -> list[str]:
    """List available samplers."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/sdapi/v1/samplers", timeout=5)
        if resp.status_code == 200:
            return [s.get("name", "") for s in resp.json()]
    except Exception:
        pass
    return ["Euler", "Euler a", "DPM++ 2M", "DPM++ SDE", "DDIM"]


def list_upscalers() -> list[str]:
    """List available upscalers."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/sdapi/v1/upscalers", timeout=5)
        if resp.status_code == 200:
            return [u.get("name", "") for u in resp.json()]
    except Exception:
        pass
    return []


def get_current_model() -> str:
    """Get the currently loaded model."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/sdapi/v1/options", timeout=5)
        if resp.status_code == 200:
            return resp.json().get("sd_model_checkpoint", "")
    except Exception:
        pass
    return ""


def switch_model(model_title: str) -> bool:
    """Switch the active SD model."""
    try:
        import httpx
        resp = httpx.post(
            f"{get_url()}/sdapi/v1/options",
            json={"sd_model_checkpoint": model_title},
            timeout=60,
        )
        return resp.status_code == 200
    except Exception:
        return False


def txt2img(
    prompt: str,
    negative_prompt: str = "ugly, blurry, low quality, deformed, disfigured",
    width: int = 1024,
    height: int = 1024,
    steps: int = 20,
    cfg_scale: float = 7.0,
    sampler: str = "Euler",
    seed: int = -1,
    batch_size: int = 1,
) -> list[bytes]:
    """Generate images from text prompt. Returns list of image bytes (PNG)."""
    import httpx

    payload = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": width,
        "height": height,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "sampler_name": sampler,
        "seed": seed,
        "batch_size": batch_size,
    }

    resp = httpx.post(
        f"{get_url()}/sdapi/v1/txt2img",
        json=payload,
        timeout=180,
    )
    resp.raise_for_status()
    data = resp.json()

    images = []
    for b64_img in data.get("images", []):
        images.append(base64.b64decode(b64_img))
    return images


def img2img(
    image_bytes: bytes,
    prompt: str,
    negative_prompt: str = "ugly, blurry, low quality",
    denoising_strength: float = 0.75,
    width: int = 1024,
    height: int = 1024,
    steps: int = 20,
    cfg_scale: float = 7.0,
    sampler: str = "Euler",
    seed: int = -1,
) -> list[bytes]:
    """Transform an existing image with a prompt. Returns list of image bytes."""
    import httpx

    b64_input = base64.b64encode(image_bytes).decode()
    payload = {
        "init_images": [b64_input],
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "denoising_strength": denoising_strength,
        "width": width,
        "height": height,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "sampler_name": sampler,
        "seed": seed,
    }

    resp = httpx.post(
        f"{get_url()}/sdapi/v1/img2img",
        json=payload,
        timeout=180,
    )
    resp.raise_for_status()
    data = resp.json()

    images = []
    for b64_img in data.get("images", []):
        images.append(base64.b64decode(b64_img))
    return images


def upscale(
    image_bytes: bytes,
    upscaler: str = "R-ESRGAN 4x+",
    scale: float = 2.0,
) -> bytes:
    """Upscale an image. Returns upscaled image bytes."""
    import httpx

    b64_input = base64.b64encode(image_bytes).decode()
    payload = {
        "resize_mode": 0,
        "upscaling_resize": scale,
        "upscaler_1": upscaler,
        "image": b64_input,
    }

    resp = httpx.post(
        f"{get_url()}/sdapi/v1/extra-single-image",
        json=payload,
        timeout=120,
    )
    resp.raise_for_status()
    return base64.b64decode(resp.json().get("image", ""))


def get_status_summary() -> str:
    """Get human-readable A1111 status."""
    if not is_available():
        return "Stable Diffusion (A1111) not detected at " + get_url()

    current = get_current_model()
    models = list_models()
    loras = list_loras()
    upscalers = list_upscalers()

    model_names = [m["title"][:40] for m in models[:5]]
    return (
        f"✅ Stable Diffusion running at {get_url()}\n"
        f"  Active model: {current or 'none'}\n"
        f"  Models: {len(models)} ({', '.join(model_names)}{'...' if len(models) > 5 else ''})\n"
        f"  LoRAs: {len(loras)}\n"
        f"  Upscalers: {len(upscalers)}"
    )
