backend/api/seedance_generation.py

"""
KIE Seedance 2 — jobs/createTask + jobs/recordInfo.
Models: bytedance/seedance-2, bytedance/seedance-2-fast.
Falls back to Replicate (bytedance/seedance-2.0 / seedance-2.0-fast) when KIE errors.
"""

import json
import logging
import os
import asyncio
from datetime import datetime
from typing import List

import httpx
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field, field_validator

from utils.public_url import get_public_base_url

router = APIRouter()
logger = logging.getLogger(__name__)

KIE_API_BASE = "https://api.kie.ai"
REPLICATE_API_BASE = "https://api.replicate.com/v1"
REPLICATE_TASK_PREFIX = "repl_"
DEFAULT_SEEDANCE_MODEL = "bytedance/seedance-2-fast"
ALLOWED_SEEDANCE_MODELS = frozenset(
    {DEFAULT_SEEDANCE_MODEL, "bytedance/seedance-2"},
)
KIE_TO_REPLICATE_SEEDANCE_MODEL = {
    "bytedance/seedance-2": "bytedance/seedance-2.0",
    "bytedance/seedance-2-fast": "bytedance/seedance-2.0-fast",
}

# In-memory callback/SSE state for Seedance tasks
seedance_results: dict[str, dict] = {}
seedance_sse_clients: dict[str, asyncio.Queue] = {}


def _normalize_terminal_state(state: str | None) -> str:
    s = (state or "").strip().lower()
    if s in {"success", "succeeded", "completed", "complete", "done"}:
        return "success"
    if s in {"fail", "failed", "error", "cancelled", "canceled"}:
        return "fail"
    return "processing"


def _extract_result_url(parsed: dict | None) -> str | None:
    if not isinstance(parsed, dict):
        return None
    # KIE commonly returns resultUrls; keep fallbacks for shape drift.
    urls = parsed.get("resultUrls") or parsed.get("urls") or []
    if isinstance(urls, list) and urls:
        first = urls[0]
        if isinstance(first, str) and first.strip():
            return first.strip()
    direct = parsed.get("resultUrl") or parsed.get("url") or parsed.get("videoUrl")
    if isinstance(direct, str) and direct.strip():
        return direct.strip()
    return None


def _cleanup_old_seedance_results(max_age_hours: int = 24) -> None:
    cutoff = datetime.now().timestamp() - (max_age_hours * 3600)
    stale = [task_id for task_id, data in seedance_results.items() if data.get("timestamp", 0) < cutoff]
    for task_id in stale:
        del seedance_results[task_id]


def _kie_key() -> str:
    api_key = os.getenv("KIE_API_KEY")
    if not api_key:
        raise HTTPException(status_code=500, detail="KIE_API_KEY not configured on server.")
    return api_key


def _replicate_token() -> str | None:
    raw = (os.getenv("REPLICATE_API_TOKEN") or os.getenv("REPLICATE_API_KEY") or "").strip()
    return raw or None


def _http_url_for_provider(u: str, public_base: str) -> str:
    """Replicate requires http(s) URIs. Map asset:// to our public base if needed."""
    s = (u or "").strip()
    if s.startswith(("http://", "https://")):
        return s
    if s.startswith("asset://"):
        rest = s[9:].lstrip("/")
        return f"{public_base.rstrip('/')}/{rest}"
    return s


async def _kie_seedance_create_task(payload: dict, api_key: str) -> str:
    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            r = await client.post(
                f"{KIE_API_BASE}/api/v1/jobs/createTask",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
            )
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"KIE request error: {e}") from e

    if r.status_code != 200:
        try:
            err = r.json()
            msg = err.get("msg") or err.get("message") or r.text[:300]
        except (json.JSONDecodeError, ValueError):
            msg = r.text[:300]
        raise HTTPException(status_code=r.status_code, detail=f"KIE createTask failed: {msg}")

    data = r.json()
    if data.get("code") != 200:
        raise HTTPException(
            status_code=data.get("code", 502),
            detail=data.get("msg", "KIE createTask rejected"),
        )

    task_id = (data.get("data") or {}).get("taskId")
    if not task_id:
        raise HTTPException(status_code=502, detail="KIE response missing taskId")
    return task_id


def _replicate_input_from_seedance_urls(
    urls: List[str],
    prompt: str,
    aspect_ratio: str,
    duration: int,
    resolution: str,
    generate_audio: bool,
    public_base: str,
) -> dict:
    mapped = [_http_url_for_provider(u, public_base) for u in urls]
    mapped = [u for u in mapped if u.startswith(("http://", "https://"))]
    if not mapped:
        raise HTTPException(
            status_code=400,
            detail="Replicate requires public http(s) image URLs (convert or host assets first).",
        )
    n = len(mapped)
    inp: dict = {
        "prompt": prompt,
        "generate_audio": generate_audio,
        "resolution": resolution,
        "aspect_ratio": aspect_ratio,
        "duration": duration,
    }
    if n == 1:
        inp["image"] = mapped[0]
    elif n == 2:
        inp["image"] = mapped[0]
        inp["last_frame_image"] = mapped[1]
    else:
        inp["reference_images"] = mapped[:9]
    return inp


async def _replicate_seedance_start(
    model_slug: str,
    input_obj: dict,
    token: str,
) -> str:
    try:
        async with httpx.AsyncClient(timeout=90.0) as client:
            r = await client.post(
                f"{REPLICATE_API_BASE}/models/{model_slug}/predictions",
                headers={
                    "Authorization": f"Bearer {token}",
                    "Content-Type": "application/json",
                },
                json={"input": input_obj},
            )
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"Replicate request error: {e}") from e

    if r.status_code not in (200, 201):
        try:
            err = r.json()
            detail = err.get("detail") or err.get("message") or r.text[:500]
        except (json.JSONDecodeError, ValueError):
            detail = r.text[:500]
        raise HTTPException(
            status_code=502,
            detail=f"Replicate create prediction failed ({r.status_code}): {detail}",
        )

    data = r.json()
    pred_id = data.get("id")
    if not pred_id:
        raise HTTPException(status_code=502, detail="Replicate response missing prediction id")
    return f"{REPLICATE_TASK_PREFIX}{pred_id}"


async def _replicate_prediction_fetch(prediction_id: str, token: str) -> dict:
    try:
        async with httpx.AsyncClient(timeout=45.0) as client:
            r = await client.get(
                f"{REPLICATE_API_BASE}/predictions/{prediction_id}",
                headers={"Authorization": f"Bearer {token}"},
            )
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"Replicate request error: {e}") from e

    if r.status_code != 200:
        raise HTTPException(status_code=502, detail="Replicate prediction fetch failed")

    return r.json()


def _replicate_prediction_to_job_shape(body: dict) -> dict:
    status = (body.get("status") or "").lower()
    err = body.get("error")
    out_url = None
    output = body.get("output")
    if isinstance(output, str):
        out_url = output
    elif isinstance(output, list) and output:
        first = output[0]
        out_url = first if isinstance(first, str) else None

    if status == "succeeded":
        state = "success"
    elif status in {"failed", "canceled", "cancelled"}:
        state = "fail"
    else:
        state = "processing"

    fail_msg = None
    if state == "fail":
        fail_msg = err if isinstance(err, str) else (str(err) if err else "Replicate prediction failed")

    return {
        "state": state,
        "url": out_url,
        "failMsg": fail_msg,
        "failCode": None,
    }


def _clamp_seedance_duration(seconds: int) -> int:
    if seconds < 4:
        return 4
    if seconds > 15:
        return 15
    return int(seconds)


def _valid_media_ref(u: str) -> bool:
    s = str(u).strip()
    return bool(s.startswith(("http://", "https://", "asset://")))


def _normalize_seedance_prompt(raw: str) -> str:
    """OpenAPI: prompt required, 3–20000 chars."""
    p = (raw or "").strip()
    if not p:
        p = "Cinematic premium product showcase, photoreal, smooth camera motion."
    if len(p) < 3:
        p = f"{p} — product video."
    if len(p) > 20000:
        p = p[:20000]
    return p


def _seedance_input_from_urls(
    urls: List[str],
    prompt: str,
    aspect_ratio: str,
    duration: int,
    resolution: str,
    generate_audio: bool,
) -> dict:
    """
    Seedance modes are mutually exclusive (OpenAPI):
    - 1 URL → image-to-video (first frame only): first_frame_url
    - 2 URLs → first & last frames: first_frame_url + last_frame_url
    - 3+ URLs → multimodal reference-to-video: reference_image_urls only (max 9)
    Do not mix first/last/reference in one request.
    """
    n = len(urls)
    base = {
        "prompt": prompt,
        "generate_audio": generate_audio,
        "resolution": resolution,
        "aspect_ratio": aspect_ratio,
        "duration": duration,
        "nsfw_checker": False,
    }
    if n == 1:
        base["first_frame_url"] = urls[0]
    elif n == 2:
        base["first_frame_url"] = urls[0]
        base["last_frame_url"] = urls[1]
    else:
        base["reference_image_urls"] = urls[:9]
    return base


class SeedanceCreateBody(BaseModel):
    prompt: str = Field(..., max_length=20000)
    reference_image_urls: List[str] = Field(..., min_length=1)
    aspect_ratio: str = "9:16"
    duration: int = 15
    resolution: str = "480p"
    generate_audio: bool = True
    model: str = Field(default=DEFAULT_SEEDANCE_MODEL, max_length=120)

    @field_validator("resolution")
    @classmethod
    def resolution_ok(cls, v: str) -> str:
        if v not in ("480p", "720p", "1080p"):
            raise ValueError("resolution must be 480p, 720p, or 1080p")
        return v

    @field_validator("model")
    @classmethod
    def seedance_model_ok(cls, v: str) -> str:
        m = (v or "").strip()
        if m not in ALLOWED_SEEDANCE_MODELS:
            raise ValueError(
                f"model must be one of: {', '.join(sorted(ALLOWED_SEEDANCE_MODELS))}"
            )
        return m


class SeedanceCreateResponse(BaseModel):
    taskId: str
    status: str = "processing"


async def _send_seedance_sse_event(task_id: str, data: dict) -> None:
    queue = seedance_sse_clients.get(task_id)
    if queue is not None:
        await queue.put(data)


def _extract_seedance_callback_state(payload: dict) -> dict:
    """
    Normalize KIE callback payloads into a browser-friendly shape.
    Supports both:
    - jobs callback with `data.info.state/resultJson`
    - direct callback payloads with `data.state/resultJson`
    """
    data = payload.get("data") if isinstance(payload.get("data"), dict) else {}
    info = data.get("info") if isinstance(data.get("info"), dict) else {}
    src = info or data

    state = _normalize_terminal_state(src.get("state"))
    fail_msg = src.get("failMsg") or payload.get("msg")
    fail_code = src.get("failCode")
    url = None
    result_json = src.get("resultJson")
    if result_json:
        try:
            parsed = json.loads(result_json) if isinstance(result_json, str) else result_json
            url = _extract_result_url(parsed)
        except (json.JSONDecodeError, TypeError, ValueError, AttributeError):
            pass

    if state == "processing":
        code = payload.get("code")
        if code == 200 and url:
            state = "success"
        elif code not in (None, 200):
            state = "fail"
        else:
            state = "processing"

    return {
        "state": state,
        "url": url,
        "failMsg": fail_msg,
        "failCode": fail_code,
    }


@router.post("/seedance/create", response_model=SeedanceCreateResponse)
async def seedance_create(body: SeedanceCreateBody):
    """Start a Seedance 2 video task. Maps 1–2 image URLs to first/last frame I2V; 3+ to reference_image_urls."""
    dur = _clamp_seedance_duration(body.duration)

    allowed_ratio = {"16:9", "4:3", "1:1", "3:4", "9:16", "21:9", "adaptive"}
    ar = body.aspect_ratio if body.aspect_ratio in allowed_ratio else "9:16"

    urls = [u for u in body.reference_image_urls if u and _valid_media_ref(u)]
    if not urls:
        raise HTTPException(
            status_code=400,
            detail="At least one image URL (http(s) or asset://) is required.",
        )

    prompt = _normalize_seedance_prompt(body.prompt)

    public_url = get_public_base_url()
    callback_url = f"{public_url}/api/seedance/callback"

    inp = _seedance_input_from_urls(
        urls,
        prompt,
        ar,
        dur,
        body.resolution,
        body.generate_audio,
    )

    payload = {
        "model": body.model,
        "input": inp,
        "callBackUrl": callback_url,
    }

    kie_err: HTTPException | None = None
    if os.getenv("KIE_API_KEY"):
        try:
            api_key = _kie_key()
            task_id = await _kie_seedance_create_task(payload, api_key)
            return SeedanceCreateResponse(taskId=task_id)
        except HTTPException as e:
            kie_err = e
            logger.warning("KIE Seedance create failed, trying Replicate if configured: %s", e.detail)
    if not os.getenv("KIE_API_KEY"):
        logger.info("KIE_API_KEY missing; using Replicate for Seedance when token is set.")

    rep = _replicate_token()
    if not rep:
        if kie_err:
            raise kie_err
        raise HTTPException(
            status_code=500,
            detail="Neither KIE_API_KEY nor REPLICATE_API_TOKEN is configured for Seedance.",
        )

    rep_slug = KIE_TO_REPLICATE_SEEDANCE_MODEL.get(body.model)
    if not rep_slug:
        raise HTTPException(status_code=500, detail="No Replicate model mapping for this Seedance variant.")

    rep_input = _replicate_input_from_seedance_urls(
        urls,
        prompt,
        ar,
        dur,
        body.resolution,
        body.generate_audio,
        public_url,
    )
    task_id = await _replicate_seedance_start(rep_slug, rep_input, rep)
    if kie_err:
        logger.info("Seedance task started via Replicate fallback (%s)", task_id)
    return SeedanceCreateResponse(taskId=task_id)


@router.get("/seedance/status/{task_id}")
async def seedance_status(task_id: str):
    """Poll KIE jobs/recordInfo or Replicate prediction; normalize for the browser."""
    cached = seedance_results.get(task_id)
    if cached:
        return cached

    if task_id.startswith(REPLICATE_TASK_PREFIX):
        token = _replicate_token()
        if not token:
            raise HTTPException(status_code=500, detail="REPLICATE_API_TOKEN not configured on server.")
        pred_id = task_id[len(REPLICATE_TASK_PREFIX) :]
        raw = await _replicate_prediction_fetch(pred_id, token)
        out = _replicate_prediction_to_job_shape(raw)
        if out["state"] in {"success", "fail"}:
            seedance_results[task_id] = {
                **out,
                "timestamp": datetime.now().timestamp(),
            }
            _cleanup_old_seedance_results()
        return out

    api_key = _kie_key()
    try:
        async with httpx.AsyncClient(timeout=45.0) as client:
            r = await client.get(
                f"{KIE_API_BASE}/api/v1/jobs/recordInfo",
                params={"taskId": task_id},
                headers={"Authorization": f"Bearer {api_key}"},
            )
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"KIE request error: {e}") from e

    if r.status_code != 200:
        raise HTTPException(status_code=r.status_code, detail="KIE recordInfo failed")

    body = r.json()
    if body.get("code") != 200:
        raise HTTPException(
            status_code=body.get("code", 502),
            detail=body.get("msg", "KIE recordInfo error"),
        )

    d = body.get("data") or {}
    state = _normalize_terminal_state(d.get("state"))
    out: dict = {
        "state": state,
        "url": None,
        "failMsg": d.get("failMsg"),
        "failCode": d.get("failCode"),
    }

    if state == "success" and d.get("resultJson"):
        try:
            rj = json.loads(d["resultJson"])
            out["url"] = _extract_result_url(rj)
        except (json.JSONDecodeError, TypeError, KeyError):
            pass

    if out["state"] in {"success", "fail"}:
        seedance_results[task_id] = {
            **out,
            "timestamp": datetime.now().timestamp(),
        }
        _cleanup_old_seedance_results()
    return out


@router.get("/seedance/events/{task_id}")
async def seedance_events(task_id: str):
    """Server-Sent Events stream for callback-driven Seedance status."""

    async def event_generator():
        queue: asyncio.Queue = asyncio.Queue()
        seedance_sse_clients[task_id] = queue
        try:
            existing = seedance_results.get(task_id)
            if existing:
                yield f"data: {json.dumps(existing)}\n\n"
                return

            if task_id.startswith(REPLICATE_TASK_PREFIX):
                token = _replicate_token()
                if not token:
                    yield f"data: {json.dumps({'state': 'fail', 'url': None, 'failMsg': 'REPLICATE_API_TOKEN not configured'})}\n\n"
                    return
                pred_id = task_id[len(REPLICATE_TASK_PREFIX) :]
                while True:
                    raw = await _replicate_prediction_fetch(pred_id, token)
                    out = _replicate_prediction_to_job_shape(raw)
                    if out["state"] in {"success", "fail"}:
                        seedance_results[task_id] = {
                            **out,
                            "timestamp": datetime.now().timestamp(),
                        }
                        _cleanup_old_seedance_results()
                        yield f"data: {json.dumps(out)}\n\n"
                        return
                    await asyncio.sleep(2.0)

            while True:
                data = await queue.get()
                yield f"data: {json.dumps(data)}\n\n"
                if data.get("state") in {"success", "fail"}:
                    return
        except asyncio.CancelledError:
            return
        finally:
            if task_id in seedance_sse_clients:
                del seedance_sse_clients[task_id]

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        },
    )


@router.post("/seedance/callback")
async def seedance_callback(request: Request):
    """
    KIE callback endpoint; updates in-memory result and notifies SSE clients.
    """
    try:
        payload = await request.json()
        if not isinstance(payload, dict):
            payload = {}
    except Exception:
        payload = {}

    data = payload.get("data") if isinstance(payload.get("data"), dict) else {}
    info = data.get("info") if isinstance(data.get("info"), dict) else {}
    task_id = info.get("taskId") or data.get("taskId")

    if task_id:
        out = _extract_seedance_callback_state(payload)
        seedance_results[task_id] = {
            **out,
            "timestamp": datetime.now().timestamp(),
        }
        await _send_seedance_sse_event(task_id, out)
        _cleanup_old_seedance_results()

    return JSONResponse(status_code=200, content={"code": 200, "msg": "ok"})