from __future__ import annotations

import asyncio
import json
from typing import Any, Awaitable, Callable

import httpx

from models.graph import GraphNode, PortValueDict
from models.events import ExecutionEvent, ProgressEvent
from services.cancellation import schedule_detached_cancel
from services.output import get_run_dir, save_base64_image

FAL_QUEUE_BASE = "https://queue.fal.run"
STREAMING_FAL_ENDPOINTS = {"openai/gpt-image-2", "openai/gpt-image-2/edit"}


def _fal_cancel_url(submit_data: dict[str, Any], endpoint_id: str, request_id: str) -> str:
    """The cancel URL for a queued request. FAL returns ``cancel_url`` in the submit
    response — prefer it, because the canonical request URL uses the *app_id*, which for a
    sub-pathed endpoint (e.g. ``fal-ai/stable-audio-25/text-to-audio``) is NOT the full
    endpoint_id (rebuilding from endpoint_id yields a 405). Fall back to deriving from
    response_url, then a last-resort construct."""
    url = submit_data.get("cancel_url")
    if url:
        return str(url)
    base = submit_data.get("response_url") or f"{FAL_QUEUE_BASE}/{endpoint_id}/requests/{request_id}"
    return str(base).rstrip("/") + "/cancel"


async def _cancel_fal_request(cancel_url: str, api_key: str) -> None:
    """Best-effort: PUT FAL's cancel endpoint so a queued/in-flight request stops on their
    side instead of running to completion (~10 min). Swallow all errors — fire-and-forget."""
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            await client.put(cancel_url, headers={"Authorization": f"Key {api_key}"})
    except Exception:
        pass


def _schedule_fal_cancel(cancel_url: str, api_key: str) -> None:
    """Fire the provider-side cancel on a DETACHED task so it survives the cancellation
    of the node's own task (which is what triggered this in the first place)."""
    schedule_detached_cancel(lambda: _cancel_fal_request(cancel_url, api_key))


def _to_fal_url(value: str) -> str:
    """Convert a local file path to a data URI, or pass URLs through."""
    if value.startswith(("http://", "https://", "data:")):
        return value
    # Local file path — convert to data URI
    import base64
    from pathlib import Path
    p = Path(value)
    if not p.exists():
        return value  # Let FAL handle the error
    suffix = p.suffix.lstrip(".").lower()
    mime_map = {
        "png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
        "webp": "image/webp", "gif": "image/gif", "mp4": "video/mp4",
        "mp3": "audio/mpeg", "wav": "audio/wav",
    }
    mime = mime_map.get(suffix, "application/octet-stream")
    b64 = base64.b64encode(p.read_bytes()).decode()
    return f"data:{mime};base64,{b64}"


def _build_fal_stream_body(node: GraphNode, inputs: dict[str, PortValueDict]) -> dict[str, Any]:
    """Build the JSON request body for FAL streaming gpt-image-2 endpoints."""
    endpoint_id = str(node.params.get("endpoint_id", "")).strip("/")

    # Validate: edit endpoint requires at least one reference image.
    if endpoint_id == "openai/gpt-image-2/edit":
        images_val = inputs.get("images")
        valid = False
        if images_val and images_val.value:
            raw = images_val.value
            if isinstance(raw, list):
                valid = any(v for v in raw)
            else:
                valid = bool(raw)
        if not valid:
            raise ValueError("At least one reference image is required for gpt-image-2 edit")

    body: dict[str, Any] = {}

    prompt_input = inputs.get("prompt")
    if prompt_input and prompt_input.value:
        body["prompt"] = str(prompt_input.value)

    # Edit endpoint: accept multiple images via the "images" input port
    images_input = inputs.get("images")
    if images_input and images_input.value:
        raw = images_input.value
        if isinstance(raw, list):
            body["image_urls"] = [_to_fal_url(str(v)) for v in raw if v]
        else:
            body["image_urls"] = [_to_fal_url(str(raw))]

    INTERNAL_KEYS = {"endpoint_id"}
    for param_key, param_val in node.params.items():
        if param_key not in INTERNAL_KEYS and param_val is not None and param_val != "":
            body[param_key] = param_val

    return body


async def handle_fal_universal(
    node: GraphNode,
    inputs: dict[str, PortValueDict],
    api_keys: dict[str, str],
    emit: Callable[[ExecutionEvent], Awaitable[None]] | None = None,
) -> dict[str, Any]:
    api_key = api_keys.get("FAL_KEY")
    if not api_key:
        raise ValueError("FAL_KEY is required")

    endpoint_id = node.params.get("endpoint_id", "")
    if not endpoint_id:
        raise ValueError("FAL endpoint ID is required (e.g. fal-ai/flux-pro/v1.1-ultra)")

    endpoint_id = str(endpoint_id).strip("/")

    # Route gpt-image-2 endpoints through the SSE streaming path
    if endpoint_id in STREAMING_FAL_ENDPOINTS and emit is not None:
        from execution.stream_runner import StreamConfig, stream_execute_image

        request_body = _build_fal_stream_body(node, inputs)
        run_dir = get_run_dir()
        config = StreamConfig(
            url=f"{FAL_QUEUE_BASE}/{endpoint_id}/stream",
            headers={
                "Authorization": f"Key {api_key}",
                "Accept": "text/event-stream",
                "Content-Type": "application/json",
            },
            timeout=180.0,
        )
        final = await stream_execute_image(
            config=config,
            request_body=request_body,
            node_id=node.id,
            emit=emit,
            run_dir=run_dir,
            provider="fal",
        )
        return {"image": {"type": "Image", "value": final}}

    # Build input payload
    fal_input: dict[str, Any] = {}

    # Map connected inputs
    prompt_input = inputs.get("prompt")
    if prompt_input and prompt_input.value:
        fal_input["prompt"] = str(prompt_input.value)

    image_input = inputs.get("image")
    if image_input and image_input.value:
        fal_input["image_url"] = _to_fal_url(str(image_input.value))

    # Multi-image port (gpt-image-1-5-edit, seedance-2-r2v) → image_urls list
    images_input = inputs.get("images")
    if images_input and images_input.value:
        raw = images_input.value
        if isinstance(raw, list):
            urls = [_to_fal_url(str(v)) for v in raw if v]
        else:
            urls = [_to_fal_url(str(raw))]
        if urls:
            fal_input["image_urls"] = urls

    # Validate: gpt-image-1-5-edit requires at least one reference image
    endpoint_id = node.params.get("endpoint_id", "")
    if endpoint_id.startswith("fal-ai/gpt-image-1.5/edit") and not fal_input.get("image_urls"):
        raise ValueError("gpt-image-1.5 edit requires at least one reference image")

    # Video input (luma-ray2-flash-modify) → video_url
    video_input = inputs.get("video")
    if video_input and video_input.value:
        fal_input["video_url"] = _to_fal_url(str(video_input.value))

    # Audio input (ltx-2-3) → audio_url
    audio_input = inputs.get("audio")
    if audio_input and audio_input.value:
        fal_input["audio_url"] = _to_fal_url(str(audio_input.value))

    # Mask input (flux-pro/v1/fill, ideogram/v3/edit inpainting) → mask_url
    mask_input = inputs.get("mask")
    if mask_input and mask_input.value:
        fal_input["mask_url"] = _to_fal_url(str(mask_input.value))

    # Character reference images (fal-ai/ideogram/character) → reference_image_urls
    reference_images_input = inputs.get("reference_images")
    if reference_images_input and reference_images_input.value:
        raw_refs = reference_images_input.value
        ref_urls = (
            [_to_fal_url(str(v)) for v in raw_refs if v]
            if isinstance(raw_refs, list)
            else [_to_fal_url(str(raw_refs))]
        )
        if ref_urls:
            fal_input["reference_image_urls"] = ref_urls

    texture_image_input = inputs.get("texture_image")
    if texture_image_input and texture_image_input.value:
        fal_input["texture_image_url"] = _to_fal_url(str(texture_image_input.value))

    end_image_input = inputs.get("end_image")
    if end_image_input and end_image_input.value:
        fal_input["end_image_url"] = _to_fal_url(str(end_image_input.value))

    tail_image_input = inputs.get("tail_image")
    if tail_image_input and tail_image_input.value:
        fal_input["tail_image_url"] = _to_fal_url(str(tail_image_input.value))

    # Multi-image inputs for 3D models (Hunyuan3D V3 Image-to-3D)
    # Hunyuan3D uses "input_image_url" for the primary image
    front_image = inputs.get("front_image")
    if front_image and front_image.value:
        fal_input["input_image_url"] = _to_fal_url(str(front_image.value))

    back_image = inputs.get("back_image")
    if back_image and back_image.value:
        fal_input["back_image_url"] = _to_fal_url(str(back_image.value))

    left_image = inputs.get("left_image")
    if left_image and left_image.value:
        fal_input["left_image_url"] = _to_fal_url(str(left_image.value))

    right_image = inputs.get("right_image")
    if right_image and right_image.value:
        fal_input["right_image_url"] = _to_fal_url(str(right_image.value))

    # Map node params (excluding our internal keys)
    INTERNAL_KEYS = {"endpoint_id"}
    for param_key, param_val in node.params.items():
        if param_key not in INTERNAL_KEYS and param_val is not None and param_val != "":
            fal_input[param_key] = param_val

    headers = {
        "Authorization": f"Key {api_key}",
        "Content-Type": "application/json",
    }

    async def noop_emit(event: ExecutionEvent) -> None:
        pass

    _emit = emit or noop_emit

    async with httpx.AsyncClient(timeout=30.0) as client:
        # Step 1: Submit to queue
        submit_resp = await client.post(
            f"{FAL_QUEUE_BASE}/{endpoint_id}",
            headers=headers,
            json=fal_input,
        )
        if submit_resp.status_code not in (200, 201):
            raise RuntimeError(f"FAL submit failed ({submit_resp.status_code}): {submit_resp.text}")

        submit_data = submit_resp.json()
        request_id = submit_data.get("request_id")
        if not request_id:
            # Some FAL endpoints return the result directly (no queue)
            return _parse_fal_output(submit_data)

        # Step 2: Poll for status
        # Use URLs from FAL response (canonical), fall back to constructed URLs
        status_url = submit_data.get("status_url") or f"{FAL_QUEUE_BASE}/{endpoint_id}/requests/{request_id}/status"
        result_url = submit_data.get("response_url") or f"{FAL_QUEUE_BASE}/{endpoint_id}/requests/{request_id}"

        max_polls = 300
        poll_interval = 2.0

        for poll_num in range(1, max_polls + 1):
            try:
                await asyncio.sleep(poll_interval)
                status_resp = await client.get(status_url, headers=headers)
            except asyncio.CancelledError:
                # User/engine cancelled — stop the FAL job upstream, then re-raise.
                _schedule_fal_cancel(_fal_cancel_url(submit_data, endpoint_id, request_id), api_key)
                raise

            if status_resp.status_code not in (200, 202):
                raise RuntimeError(f"FAL status poll failed ({status_resp.status_code}): {status_resp.text}")

            status_data = status_resp.json()
            status = status_data.get("status", "")

            progress = min(poll_num / max_polls, 0.99)
            await _emit(ProgressEvent(node_id=node.id, value=progress))

            if status == "COMPLETED":
                break
            elif status in ("FAILED", "CANCELLED"):
                error_msg = status_data.get("error", f"FAL job failed with status: {status}")
                raise RuntimeError(f"FAL job failed: {error_msg}")
            # IN_QUEUE or IN_PROGRESS — keep polling
        else:
            raise RuntimeError(f"FAL job timed out after {max_polls} polls")

        # Step 3: Fetch result
        result_resp = await client.get(result_url, headers=headers)
        if result_resp.status_code != 200:
            raise RuntimeError(f"FAL result fetch failed ({result_resp.status_code}): {result_resp.text}")

        result_data = result_resp.json()

    # Parse output — FAL endpoints vary, but common patterns:
    # Image endpoints: {"images": [{"url": "...", "content_type": "image/png"}]}
    # Audio endpoints: {"audio_url": "..."}
    # Text endpoints: {"text": "..."}
    return _parse_fal_output(result_data)


def _parse_fal_output(data: dict[str, Any]) -> dict[str, Any]:
    """Parse FAL output into our standard port format."""
    # 3D mesh output — check before images since some 3D endpoints also return preview images
    # Meshy pattern: {"model_urls": {"glb": "url", "fbx": "url"}}
    # Hunyuan3D also populates model_urls but glb is a File dict {"url": "...", ...}
    model_urls = data.get("model_urls", {})
    if isinstance(model_urls, dict) and model_urls.get("glb"):
        glb_entry = model_urls["glb"]
        glb_url = glb_entry["url"] if isinstance(glb_entry, dict) else glb_entry
        return {"mesh": {"type": "Mesh", "value": glb_url}}

    # Hunyuan/generic pattern: {"glb": {"url": "..."}} or {"glb": "url"}
    glb = data.get("glb")
    if isinstance(glb, dict) and glb.get("url"):
        return {"mesh": {"type": "Mesh", "value": glb["url"]}}
    if isinstance(glb, str) and glb:
        return {"mesh": {"type": "Mesh", "value": glb}}

    # model_glb pattern: {"model_glb": {"url": "..."}} or {"model_glb": "url"}
    model_glb = data.get("model_glb")
    if isinstance(model_glb, dict) and model_glb.get("url"):
        return {"mesh": {"type": "Mesh", "value": model_glb["url"]}}
    if isinstance(model_glb, str) and model_glb:
        return {"mesh": {"type": "Mesh", "value": model_glb}}

    # model_mesh pattern: {"model_mesh": {"url": "..."}}
    model_mesh = data.get("model_mesh")
    if isinstance(model_mesh, dict) and model_mesh.get("url"):
        return {"mesh": {"type": "Mesh", "value": model_mesh["url"]}}
    if isinstance(model_mesh, str) and model_mesh:
        return {"mesh": {"type": "Mesh", "value": model_mesh}}

    # Image / SVG output (most common)
    # SVG check: Recraft V4 text-to-vector returns images[0] with content_type "image/svg+xml"
    images = data.get("images", [])
    if images and isinstance(images, list) and len(images) > 0:
        first_image = images[0]
        if isinstance(first_image, dict):
            url = first_image.get("url", "")
            content_type = first_image.get("content_type", "")
        else:
            url = str(first_image)
            content_type = ""
        if url:
            if "svg" in content_type.lower():
                return {"svg": {"type": "SVG", "value": url}}
            return {"image": {"type": "Image", "value": url}}

    # Single image URL
    image_url = data.get("image", {})
    if isinstance(image_url, dict) and image_url.get("url"):
        return {"image": {"type": "Image", "value": image_url["url"]}}
    if isinstance(image_url, str) and image_url:
        return {"image": {"type": "Image", "value": image_url}}

    # Audio output
    audio_url = data.get("audio_url") or data.get("audio", {}).get("url", "")
    if audio_url:
        return {"audio": {"type": "Audio", "value": audio_url}}

    # Video output
    video_url = data.get("video", {}).get("url", "") or data.get("video_url", "")
    if video_url:
        return {"video": {"type": "Video", "value": video_url}}

    # Text fallback
    text = data.get("text", data.get("output", ""))
    if text:
        return {"text": {"type": "Text", "value": str(text)}}

    # Last resort — return the raw JSON as text
    return {"text": {"type": "Text", "value": json.dumps(data, indent=2)}}


def _parse_demucs_output(data: dict[str, Any]) -> dict[str, Any]:
    """Map a Demucs stem-separation response to per-stem Audio output ports.

    Unlike most FAL endpoints (one output), Demucs returns a top-level key per stem
    (vocals/drums/bass/other, plus guitar/piano on 6-stem models), each a File dict
    {"url": ...} or a bare URL string. Each present stem becomes its own Audio port.
    """
    out: dict[str, Any] = {}
    for stem in ("vocals", "drums", "bass", "other", "guitar", "piano"):
        entry = data.get(stem)
        if isinstance(entry, dict):
            url = entry.get("url")
        elif isinstance(entry, str):
            url = entry
        else:
            url = None
        if url:
            out[stem] = {"type": "Audio", "value": url}
    if not out:
        raise RuntimeError(f"Demucs returned no recognizable stems (keys: {list(data.keys())})")
    return out


async def _fal_queue_run(
    endpoint_id: str,
    fal_input: dict[str, Any],
    api_key: str,
    emit: Callable[[ExecutionEvent], Awaitable[None]] | None,
    node_id: str,
) -> dict[str, Any]:
    """Submit a job to the FAL queue, poll to completion, and return the raw result JSON.

    Shared queue plumbing for FAL handlers that need custom output parsing (e.g. Demucs'
    multi-stem response). ``handle_fal_universal`` keeps its own inline copy; consolidating
    the two is a deliberate future cleanup, kept separate here to avoid touching the
    universal handler that 39 nodes depend on.
    """
    headers = {"Authorization": f"Key {api_key}", "Content-Type": "application/json"}

    async def _noop(event: ExecutionEvent) -> None:
        pass

    _emit = emit or _noop

    async with httpx.AsyncClient(timeout=30.0) as client:
        submit_resp = await client.post(f"{FAL_QUEUE_BASE}/{endpoint_id}", headers=headers, json=fal_input)
        if submit_resp.status_code not in (200, 201):
            raise RuntimeError(f"FAL submit failed ({submit_resp.status_code}): {submit_resp.text}")

        submit_data = submit_resp.json()
        request_id = submit_data.get("request_id")
        if not request_id:
            return submit_data

        status_url = submit_data.get("status_url") or f"{FAL_QUEUE_BASE}/{endpoint_id}/requests/{request_id}/status"
        result_url = submit_data.get("response_url") or f"{FAL_QUEUE_BASE}/{endpoint_id}/requests/{request_id}"

        max_polls = 300
        for poll_num in range(1, max_polls + 1):
            try:
                await asyncio.sleep(2.0)
                status_resp = await client.get(status_url, headers=headers)
            except asyncio.CancelledError:
                # User/engine cancelled — stop the FAL job upstream, then re-raise.
                _schedule_fal_cancel(_fal_cancel_url(submit_data, endpoint_id, request_id), api_key)
                raise
            if status_resp.status_code not in (200, 202):
                raise RuntimeError(f"FAL status poll failed ({status_resp.status_code}): {status_resp.text}")
            status_data = status_resp.json()
            status = status_data.get("status", "")
            await _emit(ProgressEvent(node_id=node_id, value=min(poll_num / max_polls, 0.99)))
            if status == "COMPLETED":
                break
            if status in ("FAILED", "CANCELLED"):
                raise RuntimeError(f"FAL job failed: {status_data.get('error', status)}")
        else:
            raise RuntimeError(f"FAL job timed out after {max_polls} polls")

        result_resp = await client.get(result_url, headers=headers)
        if result_resp.status_code != 200:
            raise RuntimeError(f"FAL result fetch failed ({result_resp.status_code}): {result_resp.text}")
        return result_resp.json()


async def handle_demucs(
    node: GraphNode,
    inputs: dict[str, PortValueDict],
    api_keys: dict[str, str],
    emit: Callable[[ExecutionEvent], Awaitable[None]] | None = None,
) -> dict[str, Any]:
    """Demucs stem separation (fal-ai/demucs): one audio input -> separate stem outputs.

    Needs a dedicated handler because the response carries per-stem top-level keys, not
    the single-output shape ``_parse_fal_output`` expects.
    """
    api_key = api_keys.get("FAL_KEY")
    if not api_key:
        raise ValueError("FAL_KEY is required")

    audio_input = inputs.get("audio")
    if not audio_input or not audio_input.value:
        raise ValueError("Audio input is required for stem separation")

    fal_input: dict[str, Any] = {"audio_url": _to_fal_url(str(audio_input.value))}
    for key, val in node.params.items():
        if key != "endpoint_id" and val is not None and val != "":
            fal_input[key] = val

    # The node exposes exactly these 4 stem output ports, and the 4-stem models
    # (htdemucs_ft etc.) return HTTP 422 if the request also asks for guitar/piano,
    # so pin the requested stems to the supported set unless explicitly overridden.
    fal_input.setdefault("stems", ["vocals", "drums", "bass", "other"])

    result = await _fal_queue_run("fal-ai/demucs", fal_input, api_key, emit, node.id)
    return _parse_demucs_output(result)
