from __future__ import annotations

import json
import time
from typing import Annotated

from mcp.server.fastmcp import FastMCP

from ..api_client import api_post_json
from ..auth import resolve_auth
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.url import pre_validate_url

TOOL = "upload_from_url"


def register(mcp: FastMCP) -> None:
    @mcp.tool(
        name="upload_from_url",
        description=(
            "Low-level upload from URL. 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. Private/internal URLs are blocked. "
            "Use only when you need explicit control over the upload lifecycle."
        ),
    )
    async def upload_from_url(
        url: Annotated[str, "Public HTTPS URL of the file to upload"],
        fileName: Annotated[str | None, "Optional filename override (derived from URL 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 {"url": url, "fileName": fileName, "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)

        url_ok, url_err = pre_validate_url(url)
        if not url_ok:
            audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "error", "errorCode": "validation", "argsDigest": sanitize_args(TOOL, args)})
            raise ValueError(url_err)

        body: dict[str, object] = {"url": url}
        if fileName:
            body["fileName"] = 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_json("/v1/agent/upload-url", auth.agent_key, body, extra_headers)  # type: ignore[arg-type]
        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})
