from __future__ import annotations

import asyncio
import base64
import sys
from pathlib import Path
from typing import Any, Awaitable, Callable
from uuid import uuid4

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

VEO_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models"


def _log(msg: str) -> None:
    print(f"[veo] {msg}", file=sys.stderr, flush=True)


async def _cancel_veo_operation(op_name: str, api_key: str) -> None:
    """Best-effort upstream cancel via the standard google.longrunning Operations
    interface (POST {op_name}:cancel). The server may return UNIMPLEMENTED."""
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            await client.post(
                f"https://generativelanguage.googleapis.com/v1beta/{op_name}:cancel",
                headers={"x-goog-api-key": api_key},
                json={},
            )
        _log(f"requested upstream cancel for {op_name}")
    except Exception:
        pass


async def _image_to_veo_payload(img_str: str) -> dict[str, Any]:
    """Convert an image path/URL/data URI to Veo's inline data format."""
    if img_str.startswith(("http://", "https://")):
        async with httpx.AsyncClient(timeout=60.0) as dl_client:
            resp = await dl_client.get(img_str)
            resp.raise_for_status()
            b64_data = base64.b64encode(resp.content).decode("ascii")
            return {"bytesBase64Encoded": b64_data, "mimeType": "image/png"}
    elif img_str.startswith("data:"):
        header, b64_data = img_str.split(",", 1)
        mime_type = header.split(":")[1].split(";")[0]
        return {"bytesBase64Encoded": b64_data, "mimeType": mime_type}
    else:
        img_path = Path(img_str)
        if img_path.exists():
            b64_data = base64.b64encode(img_path.read_bytes()).decode("ascii")
            suffix = img_path.suffix.lstrip(".").lower()
            mime_map = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", "webp": "image/webp"}
            return {"bytesBase64Encoded": b64_data, "mimeType": mime_map.get(suffix, "image/png")}
    raise ValueError(f"Image not found: {img_str}")


# Veo video extension (https://ai.google.dev/gemini-api/docs/video) continues an existing
# Veo-generated clip. Works on the 3.0/3.1 generate + fast models (NOT veo-2.0, NOT 3.1-lite).
# NOTE: Veo's reference-image ("ingredients to video") mode is documented but the live Gemini
# API rejects it ("Your use case is currently not supported") on BOTH veo-3.1-generate-preview
# and veo-3.1-fast-generate-preview with an AI Studio GOOGLE_API_KEY (verified live 2026-06-08),
# so reference images are not exposed. Re-add a `referenceImages` instance field (bytesBase64Encoded
# shape, referenceType "asset") if/when Google enables the use case on the public API.
VEO_EXTEND_MODELS = {
    "veo-3.1-generate-preview",
    "veo-3.1-fast-generate-preview",
    "veo-3.0-generate-001",
    "veo-3.0-fast-generate-001",
}


async def handle_veo(
    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("GOOGLE_API_KEY")
    if not api_key:
        raise ValueError("GOOGLE_API_KEY is required")

    prompt_input = inputs.get("prompt")
    prompt_text = str(prompt_input.value) if prompt_input and prompt_input.value else ""

    model = node.params.get("model", "veo-3.1-generate-preview")

    has_image = bool(inputs.get("image") and inputs["image"].value)
    has_last = bool(inputs.get("last_frame") and inputs["last_frame"].value)
    has_video = bool(inputs.get("video") and inputs["video"].value)

    # Capability + mutual-exclusion guards — run before any byte/network work so a
    # misconfigured node fails fast and cheaply (no generation spend).
    if has_video:
        if has_image or has_last:
            raise ValueError(
                "Veo video extension cannot be combined with a first or last frame"
            )
        if model not in VEO_EXTEND_MODELS:
            raise ValueError(f"Video extension requires a Veo 3.0/3.1 model (got '{model}')")

    instance: dict[str, Any] = {}
    if prompt_text:
        instance["prompt"] = prompt_text

    # First frame (image-to-video) — Vertex-style bytesBase64Encoded payload.
    if has_image:
        instance["image"] = await _image_to_veo_payload(str(inputs["image"].value))

    # Last frame (interpolation).
    if has_last:
        instance["lastFrame"] = await _image_to_veo_payload(str(inputs["last_frame"].value))

    # Video extension — continue an existing clip. Veo only accepts the SOURCE URI of a clip
    # it generated ("must be a video that was generated by VEO that has been processed" —
    # re-uploaded/external bytes are rejected, verified live 2026-06-08). So the input must be
    # a Gemini ``files/...`` URI; wire the upstream Veo node's "Source URI" output into this
    # node's "Extend Video" input.
    if has_video:
        video_ref = str(inputs["video"].value)
        if "generativelanguage.googleapis.com" in video_ref and "/files/" in video_ref:
            instance["video"] = {"uri": video_ref}
        else:
            raise ValueError(
                "Veo extension needs a Veo-generated clip's source URI (a Gemini files/... URI). "
                "Wire the upstream Veo node's 'Source URI' output into 'Extend Video' — Veo "
                "rejects re-uploaded or external video files."
            )

    parameters: dict[str, Any] = {}
    aspect_ratio = node.params.get("aspectRatio")
    if aspect_ratio:
        parameters["aspectRatio"] = str(aspect_ratio)
    duration = node.params.get("duration")
    if duration:
        parameters["durationSeconds"] = int(str(duration).replace("s", ""))
    resolution = node.params.get("resolution")
    if resolution:
        parameters["resolution"] = str(resolution)
    person_gen = node.params.get("personGeneration")
    if person_gen:
        parameters["personGeneration"] = str(person_gen)
    seed = node.params.get("seed")
    if seed is not None and seed != "":
        parameters["seed"] = int(seed)

    # Veo extension output is locked to 720p by the API — override any higher request.
    if has_video:
        parameters["resolution"] = "720p"

    request_body: dict[str, Any] = {
        "instances": [instance],
    }
    if parameters:
        request_body["parameters"] = parameters

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

    _emit = emit or noop_emit

    _log(f"submitting to {model}")
    auth_headers = {
        "x-goog-api-key": api_key,
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=120.0) as client:
        # Submit long-running operation
        url = f"{VEO_BASE_URL}/{model}:predictLongRunning"
        resp = await client.post(url, json=request_body, headers=auth_headers)
        _log(f"submit response: {resp.status_code}")
        if resp.status_code != 200:
            raise RuntimeError(f"Veo submit failed ({resp.status_code}): {resp.text}")

        op_data = resp.json()
        op_name = op_data.get("name")
        if not op_name:
            raise RuntimeError(f"Veo did not return operation name: {op_data}")

        _log(f"polling operation {op_name}")

        # Poll for completion
        poll_url = f"https://generativelanguage.googleapis.com/v1beta/{op_name}"
        max_polls = 300
        poll_interval = 3.0

        try:
            return await _poll_veo_operation(
                client, node, poll_url, api_key, max_polls, poll_interval, _emit
            )
        except asyncio.CancelledError:
            # User/engine cancelled — fire a best-effort upstream cancel on a
            # detached task (this task's client is being torn down), then re-raise.
            schedule_detached_cancel(lambda: _cancel_veo_operation(op_name, api_key))
            raise


async def _poll_veo_operation(
    client: httpx.AsyncClient,
    node: GraphNode,
    poll_url: str,
    api_key: str,
    max_polls: int,
    poll_interval: float,
    _emit: Callable[[ExecutionEvent], Awaitable[None]],
) -> dict[str, Any]:
    for poll_num in range(1, max_polls + 1):
        await asyncio.sleep(poll_interval)

        poll_resp = await client.get(poll_url, headers={"x-goog-api-key": api_key})
        if poll_resp.status_code != 200:
            _log(f"poll FAILED: {poll_resp.status_code}")
            raise RuntimeError(f"Veo poll failed ({poll_resp.status_code}): {poll_resp.text}")

        poll_data = poll_resp.json()
        done = poll_data.get("done", False)

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

        if poll_num % 10 == 1:
            _log(f"poll #{poll_num}: done={done}")

        if done:
            response = poll_data.get("response", {})
            videos = response.get("generateVideoResponse", {}).get("generatedSamples", [])
            if not videos:
                # Try alternate response format
                videos = response.get("generatedVideos", [])

            if videos:
                video_info = videos[0]
                # Video might be a file reference or direct data
                video_uri = video_info.get("video", {}).get("uri", "")
                if video_uri:
                    # Download the video
                    _log(f"downloading video from {video_uri[:80]}...")
                    run_dir = get_run_dir()
                    filename = f"{uuid4().hex[:12]}.mp4"
                    file_path = run_dir / filename
                    dl_resp = await client.get(
                        video_uri,
                        headers={"x-goog-api-key": api_key},
                        timeout=120.0,
                        follow_redirects=True,
                    )
                    dl_resp.raise_for_status()
                    file_path.write_bytes(dl_resp.content)
                    _log(f"saved to {file_path}")
                    # Expose the Veo source URI so a downstream Veo node can extend this
                    # clip (extension only accepts a Veo-generated files/... URI, not the
                    # downloaded file). The URI is valid for ~2 days.
                    return {
                        "video": {"type": "Video", "value": str(file_path)},
                        "source_uri": {"type": "Video", "value": video_uri},
                    }

            # Fallback: check for error
            error = poll_data.get("error")
            if error:
                raise RuntimeError(f"Veo failed: {error}")

            raise RuntimeError(f"Veo completed but no video found: {list(poll_data.keys())}")

    _log("TIMED OUT")
    raise RuntimeError(f"Veo timed out after {max_polls} polls")
