from __future__ import annotations

import argparse
import base64
from dataclasses import dataclass
import mimetypes
import os
from pathlib import Path

import httpx
from mcp.server.fastmcp import FastMCP
from starlette.responses import PlainTextResponse


JsonScalar = str | int | float | bool | None
JsonValue = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]

try:
    from dotenv import load_dotenv

    _ = load_dotenv()
except Exception:
    pass


def _first_env(*keys: str) -> str | None:
    for key in keys:
        value = os.getenv(key)
        if value:
            return value
    return None


DEFAULT_BASE_URL = (
    _first_env("seedream_base_url", "SEEDREAM_BASE_URL", "ARK_BASE_URL")
    or "https://ark.cn-beijing.volces.com/api/v3"
)
DEFAULT_MODEL = _first_env("model", "MODEL", "SEEDREAM_MODEL") or "doubao-seedream-4.0"
DEFAULT_API_KEY = _first_env(
    "apikey",
    "APIKEY",
    "SEEDREAM_APIKEY",
    "SEEDREAM_API_KEY",
    "ARK_API_KEY",
)

mcp = FastMCP(
    name="SeedreamMCP",
    instructions="Seedream image generation tools over MCP",
)

app = mcp.streamable_http_app()


def _mask_secret(secret: str | None) -> str:
    if not secret:
        return ""
    if len(secret) <= 4:
        return "*" * len(secret)
    return f"{'*' * (len(secret) - 4)}{secret[-4:]}"


def _resolve_model(model: str | None) -> str:
    return (model or DEFAULT_MODEL).strip()


def _resolve_api_key(apikey: str | None) -> str:
    key = (apikey or DEFAULT_API_KEY or "").strip()
    if not key:
        raise ValueError(
            "Missing API key. Set env `apikey`/`APIKEY`/`SEEDREAM_API_KEY` or pass apikey in tool call."
        )
    return key


def _is_url_or_data_uri(value: str) -> bool:
    return value.startswith("http://") or value.startswith("https://") or value.startswith("data:image/")


def _as_data_uri(path: Path) -> str:
    mime_type, _ = mimetypes.guess_type(path.name)
    if not mime_type:
        mime_type = "image/png"
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:{mime_type};base64,{encoded}"


def _is_probable_base64(raw: str) -> bool:
    candidate = "".join(raw.split())
    if len(candidate) < 64:
        return False
    try:
        _ = base64.b64decode(candidate, validate=True)
    except Exception:
        return False
    return True


def _normalize_image_input(image: str) -> str:
    value = image.strip()
    if _is_url_or_data_uri(value):
        return value

    candidate = Path(value)
    if candidate.is_file():
        return _as_data_uri(candidate)

    if _is_probable_base64(value):
        return f"data:image/png;base64,{''.join(value.split())}"

    raise ValueError(
        "Invalid `image`. Provide an HTTP/HTTPS URL, data URI, local image path, or raw base64 string."
    )


async def _call_images_api(
    payload: dict[str, JsonValue],
    apikey: str,
    timeout_seconds: int,
    base_url: str | None = None,
) -> dict[str, object]:
    effective_base_url = (base_url or DEFAULT_BASE_URL).rstrip("/")
    endpoint = f"{effective_base_url}/images/generations"
    headers = {
        "Authorization": f"Bearer {apikey}",
        "Content-Type": "application/json",
    }

    async with httpx.AsyncClient(timeout=float(timeout_seconds)) as client:
        response = await client.post(endpoint, headers=headers, json=payload)

    try:
        data: object = response.json()
    except Exception:
        data = {"raw_text": response.text}

    if response.status_code >= 400:
        error_value: object = data
        if isinstance(data, dict) and "error" in data:
            error_value = data["error"]

        return {
            "success": False,
            "status_code": response.status_code,
            "base_url": effective_base_url,
            "error": error_value,
            "response": data,
        }

    return {
        "success": True,
        "status_code": response.status_code,
        "base_url": effective_base_url,
        "response": data,
    }


@mcp.custom_route("/health", methods=["GET"])
async def health_check(_request: object | None = None):
    return PlainTextResponse("OK")


@mcp.tool()
def get_seedream_config(
    model: str | None = None,
    apikey: str | None = None,
    base_url: str | None = None,
) -> dict[str, object]:
    effective_model = _resolve_model(model)
    effective_apikey = (apikey or DEFAULT_API_KEY or "").strip()
    effective_base_url = (base_url or DEFAULT_BASE_URL).rstrip("/")
    return {
        "base_url": effective_base_url,
        "model": effective_model,
        "apikey_configured": bool(effective_apikey),
        "apikey_masked": _mask_secret(effective_apikey),
    }


@mcp.tool()
async def seedream_text_to_image(
    prompt: str,
    model: str | None = None,
    apikey: str | None = None,
    base_url: str | None = None,
    size: str = "1024x1024",
    response_format: str = "url",
    watermark: bool = True,
    sequential_image_generation: str = "disabled",
    sequential_max_images: int | None = None,
    seed: int | None = None,
    timeout_seconds: int = 120,
) -> dict[str, object]:
    try:
        payload: dict[str, JsonValue] = {
            "model": _resolve_model(model),
            "prompt": prompt,
            "size": size,
            "response_format": response_format,
            "watermark": watermark,
            "sequential_image_generation": sequential_image_generation,
        }
        if seed is not None:
            payload["seed"] = seed
        if sequential_max_images is not None:
            payload["sequential_image_generation_options"] = {
                "max_images": sequential_max_images
            }

        return await _call_images_api(
            payload=payload,
            apikey=_resolve_api_key(apikey),
            timeout_seconds=timeout_seconds,
            base_url=base_url,
        )
    except Exception as exc:
        return {"success": False, "error": str(exc)}


@mcp.tool()
async def seedream_image_to_image(
    prompt: str,
    image: str,
    model: str | None = None,
    apikey: str | None = None,
    base_url: str | None = None,
    size: str = "adaptive",
    response_format: str = "url",
    watermark: bool = True,
    seed: int | None = None,
    timeout_seconds: int = 120,
) -> dict[str, object]:
    try:
        payload: dict[str, JsonValue] = {
            "model": _resolve_model(model),
            "prompt": prompt,
            "image": _normalize_image_input(image),
            "size": size,
            "response_format": response_format,
            "watermark": watermark,
        }
        if seed is not None:
            payload["seed"] = seed

        return await _call_images_api(
            payload=payload,
            apikey=_resolve_api_key(apikey),
            timeout_seconds=timeout_seconds,
            base_url=base_url,
        )
    except Exception as exc:
        return {"success": False, "error": str(exc)}


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Seedream MCP server launcher")
    _ = parser.add_argument(
        "--transport",
        choices=["stdio", "http", "sse"],
        default=os.getenv("SEEDREAM_MCP_TRANSPORT", os.getenv("MCP_TRANSPORT", "stdio")),
    )
    _ = parser.add_argument(
        "--host",
        default=os.getenv("SEEDREAM_MCP_HOST", os.getenv("MCP_SERVER_HOST", "127.0.0.1")),
    )
    _ = parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("SEEDREAM_MCP_PORT", os.getenv("MCP_SERVER_PORT", "8004"))),
    )
    return parser


@dataclass(frozen=True)
class LaunchArgs:
    transport: str
    host: str
    port: int


def _parse_launch_args(argv: list[str] | None = None) -> LaunchArgs:
    parsed = _build_parser().parse_args(argv)
    return LaunchArgs(
        transport=str(getattr(parsed, "transport", "stdio")),
        host=str(getattr(parsed, "host", "127.0.0.1")),
        port=int(getattr(parsed, "port", "8004")),
    )


def main(argv: list[str] | None = None) -> int:
    args = _parse_launch_args(argv)

    if args.transport.lower() in {"http", "sse"}:
        import uvicorn

        uvicorn.run(app, host=args.host, port=args.port, log_level="info")
        return 0

    mcp.run()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
