#!/usr/bin/env python3
"""
Dream Weaver API Gateway v2 — Dynamic Keyword → Local LLM → ComfyUI Pipeline
========================================================================
Port: 8080 (public-facing)
ComfyUI: localhost:8188 (internal)

NEW IN v2:
  - POST /dream-weaver now accepts keywords[] + room_type for LLM-based prompt generation
  - POST /dream-weaver/expand — expand keywords to prompt WITHOUT generating (preview)
  - GET  /room-types — list available room types
  - Uses local Ollama model (qwen3.5:27b) for prompt expansion (no cloud API dependencies)

Environment variables:
  OLLAMA_URL       — Ollama server (default: http://localhost:11434)
  OLLAMA_MODEL     — Model name (default: qwen3.5:27b)
"""
import asyncio, json, time, uuid, io, sys, os, logging, traceback
from pathlib import Path
from typing import Optional, List

import httpx
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Request
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from gateway_auth import load_gateway_api_key, is_gateway_request_authorized

# Add scripts dir to path so we can import prompt_expander
SCRIPTS_DIR = Path(__file__).parent / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))

try:
    from prompt_expander import expand_prompt, ROOM_CONTEXTS, ExpandedPrompt
    LLM_AVAILABLE = True
except ImportError:
    LLM_AVAILABLE = False
    logging.warning("prompt_expander not found — LLM expansion disabled")

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("DreamWeaverGateway")

COMFY = (os.environ.get("COMFYUI_URL") or os.environ.get("COMFY_URL") or "http://127.0.0.1:8118").rstrip("/")
COMFY_TLS_VERIFY = os.environ.get("COMFYUI_TLS_VERIFY", "true").strip().lower() not in {"0", "false", "no", "off"}
COMFY_ROOT = "/opt/dlami/nvme/ComfyUI"
GATEWAY_API_KEY = load_gateway_api_key()
PREFERRED_CHECKPOINTS = [
    "realvisxlV50_v50LightningBakedvae.safetensors",
    "realvisxlV50Lightning_v50Lightning.safetensors",
]


def comfy_client(timeout: float = 30) -> httpx.AsyncClient:
    return httpx.AsyncClient(timeout=timeout, verify=COMFY_TLS_VERIFY, follow_redirects=True)


async def list_comfy_checkpoints() -> list[str]:
    async with comfy_client(timeout=10) as client:
        response = await client.get(f"{COMFY}/models/checkpoints")
        response.raise_for_status()
        payload = response.json()
    if isinstance(payload, list):
        return [item for item in payload if isinstance(item, str)]
    return []


async def resolve_checkpoint() -> str:
    checkpoints = await list_comfy_checkpoints()
    if not checkpoints:
        raise HTTPException(
            status_code=503,
            detail=(
                "ComfyUI is online but has no checkpoint models installed. "
                "Hydrate RealVisXL into ComfyUI/models/checkpoints before generating."
            ),
        )
    lower_lookup = {item.lower(): item for item in checkpoints}
    for preferred in PREFERRED_CHECKPOINTS:
        match = lower_lookup.get(preferred.lower())
        if match:
            return match
    return checkpoints[0]

app = FastAPI(
    title="Dream Weaver API v2",
    version="2.0.0",
    description="Dynamic keyword-to-interior-design generation powered by LLM + ComfyUI"
)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# In-memory job store (swap for Redis in production)
jobs: dict = {}


def ensure_gateway_auth(request: Request) -> None:
    if is_gateway_request_authorized(request.headers, GATEWAY_API_KEY):
        return
    raise HTTPException(status_code=401, detail="Dream Weaver gateway API key is required or invalid.")


# ─── Models ──────────────────────────────────────────────────────────────────
class ExpandRequest(BaseModel):
    keywords: List[str]
    room_type: str = "living_room"
    additional_notes: str = ""


class ExpandResponse(BaseModel):
    style_name: str
    positive_prompt: str
    negative_prompt: str
    cfg: float
    denoise: float
    steps: int
    reasoning: str
    source: str


# ─── ComfyUI helpers ──────────────────────────────────────────────────────────
async def upload_to_comfy(data: bytes, filename: str) -> str:
    async with comfy_client(timeout=30) as client:
        r = await client.post(f"{COMFY}/upload/image",
                              files={"image": (filename, data, "image/jpeg")},
                              data={"overwrite": "true"})
        r.raise_for_status()
        return r.json()["name"]


def build_workflow(img_name: str, expanded: "ExpandedPrompt", ckpt_name: str) -> dict:
    """Build ComfyUI API workflow from an ExpandedPrompt result."""
    return {
        "1": {"class_type": "CheckpointLoaderSimple",
              "inputs": {"ckpt_name": ckpt_name}},
        "2": {"class_type": "LoadImage",
              "inputs": {"image": img_name, "upload": "image"}},
        "3": {"class_type": "CLIPTextEncode",   # Positive prompt
              "inputs": {"text": expanded.positive_prompt, "clip": ["1", 1]}},
        "4": {"class_type": "CLIPTextEncode",   # Negative prompt
              "inputs": {"text": expanded.negative_prompt, "clip": ["1", 1]}},
        "5": {"class_type": "VAEEncode",
              "inputs": {"pixels": ["2", 0], "vae": ["1", 2]}},
        "6": {"class_type": "KSampler",
              "inputs": {"model": ["1", 0],
                         "positive": ["3", 0],
                         "negative": ["4", 0],
                         "latent_image": ["5", 0],
                         "seed": int(time.time()) % 999983,
                         "steps": expanded.steps,
                         "cfg": expanded.cfg,
                         "sampler_name": "dpmpp_2m",
                         "scheduler": "karras",
                         "denoise": expanded.denoise}},
        "7": {"class_type": "VAEDecode",
              "inputs": {"samples": ["6", 0], "vae": ["1", 2]}},
        "8": {"class_type": "SaveImage",
              "inputs": {"images": ["7", 0],
                         "filename_prefix": f"dw_{expanded.style_name.replace(' ', '_')[:30]}"}},
    }


async def queue_prompt(workflow: dict) -> str:
    async with comfy_client(timeout=30) as client:
        r = await client.post(f"{COMFY}/prompt",
                              json={"prompt": workflow, "client_id": str(uuid.uuid4())})
        if r.status_code >= 400:
            detail = r.text
            try:
                detail = json.dumps(r.json())
            except Exception:
                pass
            raise HTTPException(status_code=502, detail=f"ComfyUI rejected Dream Weaver workflow: {detail}")
        return r.json()["prompt_id"]


async def poll_result(prompt_id: str, timeout: int = 300):
    start = time.time()
    async with comfy_client(timeout=10) as client:
        while time.time() - start < timeout:
            r = await client.get(f"{COMFY}/history/{prompt_id}")
            if r.status_code == 200:
                h = r.json().get(prompt_id, {})
                if h.get("status", {}).get("status_str") == "error":
                    return None, h.get("status", {}).get("messages", ["unknown"])
                imgs = [img for nd in h.get("outputs", {}).values()
                        for img in nd.get("images", [])]
                if imgs:
                    return imgs[0], None
            await asyncio.sleep(2)
    return None, "timeout"


async def background_poll(job_id: str, prompt_id: str):
    img, err = await poll_result(prompt_id)
    if img:
        jobs[job_id].update({"status": "done", "output": img, "completed": time.time()})
    else:
        jobs[job_id].update({"status": "error", "error": str(err)})


# ─── Endpoints ───────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    comfy_ok = False
    checkpoints: list[str] = []
    try:
        async with comfy_client(timeout=5) as c:
            r = await c.get(f"{COMFY}/system_stats")
            comfy_ok = r.status_code == 200
    except Exception:
        pass
    if comfy_ok:
        try:
            checkpoints = await list_comfy_checkpoints()
        except Exception:
            checkpoints = []
    return {
        "status": "ok",
        "comfyui": comfy_ok,
        "gpu": "4x NVIDIA L4 (96GB VRAM)",
        "model": "RealVisXL V5.0 Lightning",
        "comfyui_url": COMFY,
        "checkpoint_ready": bool(checkpoints),
        "checkpoint_count": len(checkpoints),
        "preferred_checkpoints": PREFERRED_CHECKPOINTS,
        "available_checkpoints": checkpoints[:12],
        "llm_expansion": LLM_AVAILABLE,
        "version": "2.0.0",
        "auth_required": GATEWAY_API_KEY is not None,
        "auth_scheme": "x-dream-weaver-api-key"
    }


@app.get("/room-types")
async def room_types():
    """List all supported room types with their context."""
    if not LLM_AVAILABLE:
        return {"room_types": ["bedroom", "living_room", "bathroom", "kitchen",
                               "dining_room", "home_office", "hallway", "balcony"]}
    return {
        "room_types": {
            k: {
                "description": v["description"],
                "key_elements": v["key_elements"]
            }
            for k, v in ROOM_CONTEXTS.items()
        }
    }


@app.post("/dream-weaver/expand", response_model=ExpandResponse)
async def expand_endpoint(req: ExpandRequest, request: Request):
    ensure_gateway_auth(request)
    """
    Preview the LLM-generated prompt WITHOUT submitting to ComfyUI.
    Use this to let the user review/edit the prompt before generating.
    
    Request body:
    {
        "keywords": ["blue marble", "gold veins", "renaissance", "sharp contours"],
        "room_type": "bedroom",
        "additional_notes": "luxury hotel feel"
    }
    """
    if not req.keywords:
        raise HTTPException(status_code=400, detail="keywords list cannot be empty")

    try:
        if LLM_AVAILABLE:
            result = await asyncio.to_thread(
                expand_prompt,
                keywords=req.keywords,
                room_type=req.room_type,
                additional_notes=req.additional_notes
            )
        else:
            raise RuntimeError("Local LLM model is not available or disabled.")
    except Exception as e:
        logger.error(f"Prompt expansion failed: {e}")
        raise HTTPException(status_code=500, detail=f"LLM expansion failed: {str(e)}")

    return ExpandResponse(
        style_name=result.style_name,
        positive_prompt=result.positive_prompt,
        negative_prompt=result.negative_prompt,
        cfg=result.cfg,
        denoise=result.denoise,
        steps=result.steps,
        reasoning=result.reasoning,
        source=result.source
    )


@app.post("/dream-weaver")
async def dream_weaver(
    request: Request,
    image: UploadFile = File(...),
    # ── Dynamic keyword mode (new) ──
    keywords: str = Form(default=""),           # comma-separated: "blue marble, gold, renaissance"
    room_type: str = Form(default="living_room"),
    additional_notes: str = Form(default=""),
    # ── Optional overrides ──
    custom_positive: str = Form(default=""),    # skip LLM, use this prompt directly
    custom_negative: str = Form(default=""),
    denoise: float = Form(default=0.0),         # 0.0 = use LLM recommendation
    cfg_scale: float = Form(default=0.0),       # 0.0 = use LLM recommendation
):
    """
    Submit a room photo for AI redesign using dynamic keyword → LLM → ComfyUI pipeline.

    Two modes:
    1. KEYWORD MODE (recommended): Provide keywords + room_type, LLM generates prompt
    2. DIRECT MODE: Provide custom_positive + custom_negative to bypass LLM

    Returns job_id for async polling.
    """
    ensure_gateway_auth(request)
    job_id = str(uuid.uuid4())
    jobs[job_id] = {"status": "uploading", "created": time.time()}

    try:
        # Upload image to ComfyUI
        data = await image.read()
        filename = f"dw_{job_id[:8]}_{image.filename or 'room.jpg'}"
        comfy_name = await upload_to_comfy(data, filename)
        jobs[job_id]["status"] = "expanding_prompt"

        # ── Determine prompt ──────────────────────────────────────────────
        if custom_positive:
            # Direct mode — user provided prompts explicitly
            from dataclasses import dataclass

            @dataclass
            class DirectPrompt:
                style_name: str = "custom"
                positive_prompt: str = custom_positive
                negative_prompt: str = custom_negative or (
                    "(worst quality, low quality, illustration, 3d render, painting, cartoon, sketch), "
                    "blurry, distorted, deformed, extra windows, unrealistic lighting, structural changes"
                )
                cfg: float = cfg_scale or 7.5
                denoise: float = denoise or 0.72
                steps: int = 30
                reasoning: str = "Direct user input"
                source: str = "direct"

            expanded = DirectPrompt()

        elif keywords:
            # Keyword mode — expand via LLM
            kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
            if LLM_AVAILABLE:
                expanded = await asyncio.to_thread(
                    expand_prompt,
                    keywords=kw_list,
                    room_type=room_type,
                    additional_notes=additional_notes
                )
            else:
                raise HTTPException(status_code=500, detail="LLM model is not available or disabled.")

            # Apply manual overrides if provided
            if denoise > 0:
                expanded.denoise = denoise
            if cfg_scale > 0:
                expanded.cfg = cfg_scale

        else:
            raise HTTPException(status_code=400,
                                detail="Provide either 'keywords' or 'custom_positive'")

        jobs[job_id].update({
            "status": "queued",
            "style": expanded.style_name,
            "prompt_source": expanded.source,
            "positive_prompt": expanded.positive_prompt,
            "negative_prompt": expanded.negative_prompt,
            "room_type": room_type,
        })

        # Submit workflow
        ckpt_name = await resolve_checkpoint()
        jobs[job_id]["checkpoint"] = ckpt_name
        wf = build_workflow(comfy_name, expanded, ckpt_name)
        prompt_id = await queue_prompt(wf)
        jobs[job_id].update({"status": "processing", "prompt_id": prompt_id})

        # Start background polling
        asyncio.create_task(background_poll(job_id, prompt_id))

        return {
            "job_id": job_id,
            "status": "processing",
            "style": expanded.style_name,
            "prompt_preview": expanded.positive_prompt[:120] + "...",
            "reasoning": expanded.reasoning,
            "poll_url": f"/dream-weaver/status/{job_id}",
            "result_url": f"/dream-weaver/result/{job_id}"
        }

    except HTTPException:
        raise
    except Exception as e:
        jobs[job_id] = {"status": "error", "error": str(e)}
        logger.error(f"Generation failed: {e}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/dream-weaver/status/{job_id}")
async def status(job_id: str, request: Request):
    ensure_gateway_auth(request)
    job = jobs.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    result = {k: v for k, v in job.items() if k != "output"}
    result["ready"] = job.get("status") == "done"
    if result["ready"]:
        result["result_url"] = f"/dream-weaver/result/{job_id}"
    return result


@app.get("/dream-weaver/result/{job_id}")
async def result(job_id: str, request: Request):
    ensure_gateway_auth(request)
    job = jobs.get(job_id)
    if not job or job.get("status") != "done":
        raise HTTPException(status_code=404, detail="Result not ready")
    img = job["output"]
    url = (f"{COMFY}/view?filename={img['filename']}"
           f"&subfolder={img.get('subfolder','')}&type={img.get('type','output')}")
    async with comfy_client(timeout=30) as c:
        r = await c.get(url)
    return StreamingResponse(
        io.BytesIO(r.content),
        media_type="image/png",
        headers={"Content-Disposition": f"attachment; filename=dreamweaver_{job_id[:8]}.png"}
    )


@app.post("/dream-weaver/sync")
async def dream_weaver_sync(
    request: Request,
    image: UploadFile = File(...),
    keywords: str = Form(default=""),
    room_type: str = Form(default="living_room"),
    additional_notes: str = Form(default=""),
    custom_positive: str = Form(default=""),
    custom_negative: str = Form(default=""),
):
    """
    Blocking version — waits up to 120s and returns image bytes directly.
    Use for testing. Prefer async /dream-weaver for production.
    """
    ensure_gateway_auth(request)
    data = await image.read()
    filename = f"sync_{uuid.uuid4().hex[:8]}_{image.filename or 'room.jpg'}"
    comfy_name = await upload_to_comfy(data, filename)

    if custom_positive:
        from dataclasses import dataclass
        @dataclass
        class _P:
            style_name = "custom"
            positive_prompt = custom_positive
            negative_prompt = custom_negative or "(worst quality, low quality), blurry, structural changes"
            cfg = 7.5; denoise = 0.72; steps = 30
            reasoning = ""; source = "direct"
        expanded = _P()
    elif keywords:
        kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
        if LLM_AVAILABLE:
            expanded = expand_prompt(kw_list, room_type, additional_notes)
        else:
            raise RuntimeError("Local LLM model is not available or disabled.")
    else:
        raise HTTPException(status_code=400, detail="Provide keywords or custom_positive")

    ckpt_name = await resolve_checkpoint()
    wf = build_workflow(comfy_name, expanded, ckpt_name)
    prompt_id = await queue_prompt(wf)
    img, err = await poll_result(prompt_id, timeout=120)
    if err:
        raise HTTPException(status_code=500, detail=str(err))
    url = (f"{COMFY}/view?filename={img['filename']}"
           f"&subfolder={img.get('subfolder','')}&type={img.get('type','output')}")
    async with comfy_client(timeout=30) as c:
        r = await c.get(url)
    return StreamingResponse(io.BytesIO(r.content), media_type="image/png",
                             headers={"X-Style": expanded.style_name,
                                      "X-Prompt-Source": expanded.source})


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8082")), log_level="info")