# /// script
# requires-python = ">=3.10"
# dependencies = ["mcp>=1.27,<2", "httpx>=0.27,<1"]
# ///
# Generated by scripts/sync-mcp.mjs from server/app/mcp_server.py. Do not edit.
from __future__ import annotations

import argparse
import asyncio
import os
import time
from typing import Any, Literal

import httpx
from mcp.server.fastmcp import FastMCP

DEFAULT_KUBFLOW_BASE_URL = "https://kubflow.com"
DEFAULT_MCP_HOST = "127.0.0.1"
DEFAULT_MCP_PORT = 8000
DEFAULT_MCP_PATH = "/mcp"
SUPPORTED_CATEGORIES = {"image", "video"}
MCP_TRANSPORTS = {"stdio", "sse", "streamable-http"}
MCP_LOG_LEVELS = {"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"}
TERMINAL_STATUSES = {
    "succeeded",
    "success",
    "completed",
    "complete",
    "failed",
    "error",
    "cancelled",
    "canceled",
}

SERVER_INSTRUCTIONS = """
Use Kubflow to generate images and videos through the user's Kubflow API key.
Start by listing models when the user has not named a model. Create generations,
then poll status until a terminal state before returning output URLs. API calls
spend the key owner's Kubflow credits, so avoid duplicate runs unless the user
explicitly asks for another generation.
""".strip()

def _mcp_port() -> int:
    raw_port = str(os.getenv("KUBFLOW_MCP_PORT") or DEFAULT_MCP_PORT).strip()
    try:
        return int(raw_port)
    except ValueError:
        return DEFAULT_MCP_PORT


def _mcp_log_level() -> Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
    log_level = str(os.getenv("KUBFLOW_MCP_LOG_LEVEL") or "WARNING").strip().upper()
    if log_level not in MCP_LOG_LEVELS:
        return "WARNING"
    if log_level == "DEBUG":
        return "DEBUG"
    if log_level == "INFO":
        return "INFO"
    if log_level == "ERROR":
        return "ERROR"
    if log_level == "CRITICAL":
        return "CRITICAL"
    return "WARNING"


mcp = FastMCP(
    name="Kubflow",
    instructions=SERVER_INSTRUCTIONS,
    log_level=_mcp_log_level(),
    host=str(os.getenv("KUBFLOW_MCP_HOST") or DEFAULT_MCP_HOST).strip()
    or DEFAULT_MCP_HOST,
    port=_mcp_port(),
    streamable_http_path=str(os.getenv("KUBFLOW_MCP_PATH") or DEFAULT_MCP_PATH).strip()
    or DEFAULT_MCP_PATH,
)


def _normalize_transport(
    transport: str,
) -> Literal["stdio", "sse", "streamable-http"]:
    normalized = str(transport or "stdio").strip().lower().replace("_", "-")
    if normalized == "http":
        normalized = "streamable-http"
    if normalized not in MCP_TRANSPORTS:
        raise ValueError("transport must be stdio, sse, or streamable-http")
    if normalized == "sse":
        return "sse"
    if normalized == "streamable-http":
        return "streamable-http"
    return "stdio"


def _parse_transport() -> Literal["stdio", "sse", "streamable-http"]:
    parser = argparse.ArgumentParser(description="Run the Kubflow MCP server.")
    parser.add_argument(
        "--transport",
        default=os.getenv("KUBFLOW_MCP_TRANSPORT") or "stdio",
        help="MCP transport: stdio, sse, or streamable-http. Default: stdio.",
    )
    args = parser.parse_args()
    return _normalize_transport(args.transport)


def _api_base_url() -> str:
    base_url = (
        os.getenv("KUBFLOW_API_BASE_URL")
        or os.getenv("KUBFLOW_BASE_URL")
        or DEFAULT_KUBFLOW_BASE_URL
    )
    base_url = str(base_url).strip().rstrip("/")
    if not base_url:
        base_url = DEFAULT_KUBFLOW_BASE_URL
    if base_url.endswith("/api/v1"):
        return base_url
    return f"{base_url}/api/v1"


def _api_key() -> str:
    api_key = str(os.getenv("KUBFLOW_API_KEY") or "").strip()
    if not api_key:
        raise RuntimeError(
            "KUBFLOW_API_KEY is required. Create a Kubflow API key in Settings "
            "and pass it to the MCP server environment."
        )
    return api_key


def _validate_category(category: str, *, allow_empty: bool = False) -> str:
    normalized = str(category or "").strip().lower()
    if allow_empty and not normalized:
        return ""
    if normalized not in SUPPORTED_CATEGORIES:
        raise ValueError("category must be image or video")
    return normalized


def _clean_text(value: object, *, field: str, max_length: int = 8000) -> str:
    text = str(value or "").strip()
    if not text:
        raise ValueError(f"{field} is required")
    return text[:max_length]


async def _request_json(
    method: str,
    path: str,
    *,
    params: dict[str, Any] | None = None,
    json_body: dict[str, Any] | None = None,
    timeout_seconds: float = 60.0,
) -> dict[str, Any]:
    url = f"{_api_base_url()}{path}"
    headers = {
        "Authorization": f"Bearer {_api_key()}",
        "Content-Type": "application/json",
        "Accept": "application/json",
        "User-Agent": "kubflow-mcp/1.0",
    }
    async with httpx.AsyncClient(timeout=timeout_seconds) as client:
        response = await client.request(
            method,
            url,
            headers=headers,
            params=params or None,
            json=json_body,
        )

    try:
        payload = response.json()
    except ValueError:
        payload = {"body": response.text[:2000]}

    if response.status_code >= 400:
        detail = payload.get("detail") if isinstance(payload, dict) else None
        error = payload.get("error") if isinstance(payload, dict) else None
        message = detail or error
        message = str(message or response.text or "Kubflow API request failed")
        raise RuntimeError(
            f"Kubflow API {response.status_code} for {method.upper()} {path}: "
            f"{message[:1000]}"
        )

    if isinstance(payload, dict):
        return payload
    return {"result": payload}


def _compact_model(model: dict[str, Any]) -> dict[str, Any]:
    input_schema = model.get("input_schema") if isinstance(model, dict) else {}
    if not isinstance(input_schema, dict):
        input_schema = {}
    properties = input_schema.get("properties")
    if not isinstance(properties, dict):
        properties = {}
    return {
        "id": model.get("id"),
        "category": model.get("category"),
        "name": model.get("name"),
        "description": model.get("description"),
        "credits": model.get("credits"),
        "input_fields": sorted(str(key) for key in properties.keys()),
        "required_fields": input_schema.get("required") or [],
    }


def _compact_models(payload: dict[str, Any], detail: str) -> dict[str, Any]:
    if detail == "detailed":
        return payload
    models = payload.get("models")
    if not isinstance(models, list):
        return payload
    compact_models = [
        _compact_model(model) for model in models if isinstance(model, dict)
    ]
    return {"models": compact_models}


def _is_terminal_status(payload: dict[str, Any]) -> bool:
    status = str(payload.get("status") or "").strip().lower()
    return status in TERMINAL_STATUSES


@mcp.tool()
async def kubflow_account() -> dict[str, Any]:
    """Return the API key owner's Kubflow user id, plan, and credit balance."""
    return await _request_json("GET", "/me", timeout_seconds=20.0)


@mcp.tool()
async def kubflow_list_models(
    category: Literal["", "image", "video"] = "",
    detail: Literal["concise", "detailed"] = "concise",
) -> dict[str, Any]:
    """List Kubflow generation models available to the API key owner.

    Use category='image' or category='video' when the user knows the media type.
    The concise view is best for model selection. Use detailed only when input
    schemas are needed.
    """
    normalized_category = _validate_category(category, allow_empty=True)
    payload = await _request_json(
        "GET",
        "/models",
        params={"category": normalized_category} if normalized_category else None,
        timeout_seconds=30.0,
    )
    return _compact_models(payload, detail)


@mcp.tool()
async def kubflow_create_generation(
    model: str,
    prompt: str,
    category: Literal["image", "video"],
    params: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Start an asynchronous Kubflow image or video generation.

    Put model-specific options such as aspect_ratio, resolution, duration,
    image_url, video_url, or sound inside params. Returns a generation_id that
    should be passed to kubflow_get_generation until the run completes.
    """
    payload = {
        "category": _validate_category(category),
        "model": _clean_text(model, field="model", max_length=200),
        "prompt": _clean_text(prompt, field="prompt"),
        "params": params if isinstance(params, dict) else {},
    }
    return await _request_json(
        "POST",
        "/generations",
        json_body=payload,
        timeout_seconds=60.0,
    )


@mcp.tool()
async def kubflow_get_generation(generation_id: str) -> dict[str, Any]:
    """Check generation status and return result URLs or errors when available."""
    generation_key = _clean_text(generation_id, field="generation_id", max_length=200)
    return await _request_json(
        "GET",
        f"/generations/{generation_key}",
        timeout_seconds=30.0,
    )


@mcp.tool()
async def kubflow_generate_and_wait(
    model: str,
    prompt: str,
    category: Literal["image", "video"],
    params: dict[str, Any] | None = None,
    timeout_seconds: int = 180,
    poll_interval_seconds: float = 4.0,
) -> dict[str, Any]:
    """Create a generation and poll until it succeeds, fails, or times out.

    Use this for simple one-shot requests. For long video generations or agent
    workflows that need progress updates, call kubflow_create_generation first
    and then kubflow_get_generation manually.
    """
    timeout_seconds = max(10, min(int(timeout_seconds), 600))
    poll_interval_seconds = max(1.0, min(float(poll_interval_seconds), 30.0))

    generation = await kubflow_create_generation(
        model=model,
        prompt=prompt,
        category=category,
        params=params,
    )
    generation_id = str(
        generation.get("generation_id")
        or generation.get("id")
        or generation.get("run_id")
        or ""
    ).strip()
    if not generation_id:
        return {
            "status": "unknown",
            "generation": generation,
            "message": "Kubflow did not return a generation_id.",
        }

    deadline = time.monotonic() + timeout_seconds
    latest: dict[str, Any] = generation
    while time.monotonic() < deadline:
        latest = await kubflow_get_generation(generation_id)
        if _is_terminal_status(latest):
            return {"generation": generation, "latest": latest}
        await asyncio.sleep(poll_interval_seconds)

    return {
        "status": "timeout",
        "generation": generation,
        "latest": latest,
        "message": (
            "Generation is still running. Continue polling with "
            f"kubflow_get_generation using generation_id={generation_id}."
        ),
    }


def main() -> None:
    """Run the Kubflow MCP server."""
    mcp.run(transport=_parse_transport())


if __name__ == "__main__":
    main()