from __future__ import annotations

import json
import time
from typing import Annotated

from mcp.server.fastmcp import FastMCP

from ..api_client import api_post_multipart
from ..auth import resolve_auth
from ..config import get_transport_mode
from ..security.audit import audit_log, sanitize_args
from ..security.metrics import record_auth_failure, record_throttle_rejection, record_tool_call
from ..security.throttle import check_throttle
from ..validation.base64_content import validate_and_decode_base64
from ..validation.file_path import validate_and_read_file

TOOL = "upload_file"


def register(mcp: FastMCP) -> None:
    @mcp.tool(
        name="upload_file",
        description=(
            "Low-level upload. Prefer upload_and_wait instead — it handles polling automatically. "
            "This tool returns immediately with PENDING_UPLOAD status; you must poll get_upload_status yourself. "
            "Max 5 MB. Provide EITHER 'filePath' (local stdio only) OR 'content' (base64). "
            "Use only when you need explicit control over the upload lifecycle."
        ),
    )
    async def upload_file(
        filePath: Annotated[str | None, "Absolute path to a local file. Only works when this MCP server runs locally (stdio transport). Preferred over base64 — zero tokens wasted."] = None,
        content: Annotated[str | None, "Base64-encoded file content (standard or URL-safe base64)"] = None,
        fileName: Annotated[str | None, "Original filename with extension (e.g. report.pdf). Required when using 'content', optional for 'filePath'."] = None,
        mimeType: Annotated[str | None, "MIME type override (auto-detected if omitted)"] = None,
        pricePerDownload: Annotated[float | None, "Price per download in USD (e.g. 0.50). Omit for free downloads."] = None,
        idempotencyKey: Annotated[str | None, "Unique key for retry-safe uploads (8-64 alphanumeric chars)"] = None,
    ) -> str:
        start = time.monotonic()
        args = {k: v for k, v in {"filePath": filePath, "content": content, "fileName": fileName, "mimeType": mimeType, "pricePerDownload": pricePerDownload, "idempotencyKey": idempotencyKey}.items() if v is not None}

        auth = resolve_auth()
        if not auth:
            record_auth_failure()
            audit_log({"tool": TOOL, "result": "auth_failed"})
            raise ValueError("Authentication required. Set AGENTVEE_API_KEY environment variable.")

        ok, msg = check_throttle(TOOL, auth.token_prefix)
        if not ok:
            record_throttle_rejection()
            audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "throttled"})
            raise ValueError(msg)

        has_file_path = bool(filePath)
        has_content = bool(content)

        if not has_file_path and not has_content:
            raise ValueError("Provide either 'filePath' or 'content'. Neither was supplied.")
        if has_file_path and has_content:
            raise ValueError("Provide either 'filePath' or 'content', not both.")

        if has_file_path and get_transport_mode() != "stdio":
            raise ValueError("filePath is only supported when this MCP server runs locally (stdio transport). Use 'content' instead.")

        file_data: bytes
        resolved_name: str
        mime = mimeType or "application/octet-stream"

        if has_file_path:
            valid, data, name, err = validate_and_read_file(filePath, fileName)  # type: ignore[arg-type]
            if not valid:
                audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "error", "errorCode": "validation", "argsDigest": sanitize_args(TOOL, args)})
                raise ValueError(err)
            file_data = data  # type: ignore[assignment]
            resolved_name = name
        else:
            if not fileName:
                raise ValueError("'fileName' is required when uploading base64 content.")
            valid, decoded, err = validate_and_decode_base64(content)  # type: ignore[arg-type]
            if not valid:
                audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "error", "errorCode": "validation", "argsDigest": sanitize_args(TOOL, args)})
                raise ValueError(err)
            file_data = decoded  # type: ignore[assignment]
            resolved_name = fileName

        extra_headers: dict[str, str] = {}
        if idempotencyKey:
            extra_headers["Idempotency-Key"] = idempotencyKey
        if pricePerDownload is not None and pricePerDownload > 0:
            extra_headers["X-Price-Per-Download"] = str(pricePerDownload)

        result = await api_post_multipart("/v1/agent/upload", auth.agent_key, file_data, resolved_name, mime, extra_headers)
        duration_ms = (time.monotonic() - start) * 1000

        if not result.ok:
            record_tool_call(TOOL, duration_ms, True)
            audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "error", "errorCode": result.code, "durationMs": round(duration_ms), "argsDigest": sanitize_args(TOOL, args)})
            raise ValueError(result.message)

        data = result.data
        record_tool_call(TOOL, duration_ms, False)
        audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "ok", "uploadId": data.get("uploadId"), "durationMs": round(duration_ms), "argsDigest": sanitize_args(TOOL, args)})

        output = {
            "uploadId": data.get("uploadId"),
            "status": data.get("status", "PENDING_UPLOAD"),
            "downloadUrl": data.get("downloadUrl"),
            "expiresAt": data.get("expiresAt"),
        }
        return json.dumps({k: v for k, v in output.items() if v is not None})
