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

"""ComfyUI deep integration — model discovery, workflow builder, progress tracking.

Connects to a locally-running ComfyUI instance for FLUX.2, SDXL, SD 1.5,
and any installed checkpoint. Supports txt2img, img2img, and upscaling.
"""

from __future__ import annotations

import json
import logging
import os
import time
import uuid
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)

DEFAULT_URL = "http://localhost:8188"


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


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


def get_system_info() -> dict:
    """Get ComfyUI system stats (GPU, VRAM, etc.)."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/system_stats", timeout=5)
        return resp.json() if resp.status_code == 200 else {}
    except Exception:
        return {}


def list_checkpoints() -> list[str]:
    """List available checkpoint models."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/object_info/CheckpointLoaderSimple", timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            return data["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]
    except Exception:
        pass
    return []


def list_loras() -> list[str]:
    """List available LoRA models."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/object_info/LoraLoader", timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            return data["LoraLoader"]["input"]["required"]["lora_name"][0]
    except Exception:
        pass
    return []


def list_samplers() -> list[str]:
    """List available sampler methods."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/object_info/KSampler", timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            return data["KSampler"]["input"]["required"]["sampler_name"][0]
    except Exception:
        pass
    return ["euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"]


def list_schedulers() -> list[str]:
    """List available schedulers."""
    try:
        import httpx
        resp = httpx.get(f"{get_url()}/object_info/KSampler", timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            return data["KSampler"]["input"]["required"]["scheduler"][0]
    except Exception:
        pass
    return ["normal", "karras", "exponential", "sgm_uniform"]


def build_txt2img_workflow(
    prompt: str,
    negative_prompt: str = "ugly, blurry, low quality, deformed",
    checkpoint: str = "",
    width: int = 1024,
    height: int = 1024,
    steps: int = 20,
    cfg_scale: float = 7.0,
    sampler: str = "euler",
    scheduler: str = "normal",
    seed: Optional[int] = None,
    lora: Optional[str] = None,
    lora_strength: float = 0.8,
) -> dict:
    """Build a ComfyUI workflow JSON for txt2img generation."""
    if seed is None:
        seed = int(time.time()) % 2**32

    if not checkpoint:
        ckpts = list_checkpoints()
        if not ckpts:
            raise RuntimeError("No checkpoints found in ComfyUI")
        # Prefer FLUX or SDXL if available
        for preferred in ["flux", "sdxl", "sd_xl"]:
            match = [c for c in ckpts if preferred in c.lower()]
            if match:
                checkpoint = match[0]
                break
        if not checkpoint:
            checkpoint = ckpts[0]

    workflow = {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed,
                "steps": steps,
                "cfg": cfg_scale,
                "sampler_name": sampler,
                "scheduler": scheduler,
                "denoise": 1.0,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": checkpoint},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": prompt, "clip": ["4", 1]},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative_prompt, "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]},
        },
    }

    # Add LoRA if specified
    if lora:
        workflow["10"] = {
            "class_type": "LoraLoader",
            "inputs": {
                "lora_name": lora,
                "strength_model": lora_strength,
                "strength_clip": lora_strength,
                "model": ["4", 0],
                "clip": ["4", 1],
            },
        }
        # Redirect KSampler and CLIP to use LoRA output
        workflow["3"]["inputs"]["model"] = ["10", 0]
        workflow["6"]["inputs"]["clip"] = ["10", 1]
        workflow["7"]["inputs"]["clip"] = ["10", 1]

    return workflow


def generate(
    prompt: str,
    negative_prompt: str = "ugly, blurry, low quality, deformed",
    checkpoint: str = "",
    width: int = 1024,
    height: int = 1024,
    steps: int = 20,
    cfg_scale: float = 7.0,
    sampler: str = "euler",
    scheduler: str = "normal",
    seed: Optional[int] = None,
    lora: Optional[str] = None,
    lora_strength: float = 0.8,
    timeout: int = 180,
) -> bytes:
    """Generate an image using ComfyUI. Returns image bytes."""
    import httpx

    url = get_url()
    workflow = build_txt2img_workflow(
        prompt=prompt, negative_prompt=negative_prompt,
        checkpoint=checkpoint, width=width, height=height,
        steps=steps, cfg_scale=cfg_scale, sampler=sampler,
        scheduler=scheduler, seed=seed, lora=lora, lora_strength=lora_strength,
    )

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

    # Poll for completion
    start = time.time()
    while time.time() - start < timeout:
        time.sleep(1)
        try:
            history = httpx.get(f"{url}/history/{prompt_id}", timeout=5).json()
        except Exception:
            continue

        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"{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(f"ComfyUI generation timed out after {timeout}s")


def get_status_summary() -> str:
    """Get a human-readable ComfyUI status for the connect wizard."""
    if not is_available():
        return "ComfyUI not detected at " + get_url()

    info = get_system_info()
    ckpts = list_checkpoints()
    loras = list_loras()

    devices = info.get("devices", [])
    gpu_info = ""
    if devices:
        dev = devices[0]
        name = dev.get("name", "Unknown GPU")
        vram = dev.get("vram_total", 0)
        vram_free = dev.get("vram_free", 0)
        gpu_info = f"\n  GPU: {name} ({vram // (1024**3)}GB VRAM, {vram_free // (1024**3)}GB free)"

    return (
        f"✅ ComfyUI running at {get_url()}{gpu_info}\n"
        f"  Checkpoints: {len(ckpts)} ({', '.join(ckpts[:3])}{'...' if len(ckpts) > 3 else ''})\n"
        f"  LoRAs: {len(loras)}"
    )
