# services/comfy-agent/app/api.py
from fastapi import APIRouter, HTTPException
from .models import GenerateImageRequest, GenerateVideoRequest, JobStatus
from .jobs import JOB_STORE
from .worker import enqueue

router = APIRouter()

def _build_workflow_t2i(req: GenerateImageRequest) -> dict:
    # Basic SD 1.5 workflow
    # Node structure: CheckpointLoader -> CLIP Encode -> KSampler -> VAE Decode -> SaveImage
    return {
        "3": {
            "inputs": {
                "seed": req.seed if req.seed else 42,
                "steps": req.steps,
                "cfg": 7.0,
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0]
            },
            "class_type": "KSampler"
        },
        "4": {
            "inputs": {
                "ckpt_name": "v1-5-pruned-emaonly.safetensors"
            },
            "class_type": "CheckpointLoaderSimple"
        },
        "5": {
            "inputs": {
                "width": req.width,
                "height": req.height,
                "batch_size": 1
            },
            "class_type": "EmptyLatentImage"
        },
        "6": {
            "inputs": {
                "text": req.prompt,
                "clip": ["4", 1]
            },
            "class_type": "CLIPTextEncode"
        },
        "7": {
            "inputs": {
                "text": req.negative_prompt if req.negative_prompt else "text, watermark, blurry",
                "clip": ["4", 1]
            },
            "class_type": "CLIPTextEncode"
        },
        "8": {
            "inputs": {
                "samples": ["3", 0],
                "vae": ["4", 2]
            },
            "class_type": "VAEDecode"
        },
        "9": {
            "inputs": {
                "filename_prefix": "comfy-agent",
                "images": ["8", 0]
            },
            "class_type": "SaveImage"
        }
    }

def _build_workflow_t2v(req: GenerateVideoRequest) -> dict:
    # MVP placeholder for LTX-2 pipeline; replace with actual LTX-2 workflow.
    return {
        "1": {"class_type": "CLIPTextEncode", "inputs": {"text": req.prompt, "clip": ["2", 0]}},
        # TODO: Add complete workflow JSON for text-to-video with LTX-2
    }

@router.post("/generate/image", response_model=JobStatus)
async def generate_image(req: GenerateImageRequest):
    job = JOB_STORE.create("text-to-image")
    graph = _build_workflow_t2i(req)
    enqueue(job.job_id, "text-to-image", graph)
    return JOB_STORE.get(job.job_id)

@router.post("/generate/video", response_model=JobStatus)
async def generate_video(req: GenerateVideoRequest):
    job = JOB_STORE.create("text-to-video")
    graph = _build_workflow_t2v(req)
    enqueue(job.job_id, "text-to-video", graph)
    return JOB_STORE.get(job.job_id)

@router.get("/status/{job_id}", response_model=JobStatus)
async def status(job_id: str):
    job = JOB_STORE.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="job_not_found")
    return job

@router.get("/result/{job_id}", response_model=JobStatus)
async def result(job_id: str):
    job = JOB_STORE.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="job_not_found")
    return job