from __future__ import annotations

import json
import time
from typing import Annotated

from mcp.server.fastmcp import FastMCP

from ..api_client import api_get
from ..auth import resolve_auth
from ..security.audit import audit_log
from ..security.metrics import record_auth_failure, record_throttle_rejection, record_tool_call
from ..security.throttle import check_throttle
from ..validation.upload_id import validate_upload_id

TOOL = "get_upload_status"


def register(mcp: FastMCP) -> None:
    @mcp.tool(
        name="get_upload_status",
        description=(
            "Check the processing status of a previously uploaded file. "
            "Returns the current status (PENDING_UPLOAD, UPLOADING, READY, FAILED) and whether the file is ready for download. "
            "Rarely needed — upload_and_wait handles polling automatically. "
            "Use this only to check on uploads that timed out or for manual status inspection."
        ),
    )
    async def get_upload_status(
        uploadId: Annotated[str, "Upload ID returned from upload_file or upload_from_url (e.g. up_a1b2c3d4e5f6g7h8)"],
    ) -> str:
        start = time.monotonic()

        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)

        id_ok, id_err = validate_upload_id(uploadId)
        if not id_ok:
            audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "error", "errorCode": "validation"})
            raise ValueError(id_err)

        result = await api_get(f"/v1/upload/{uploadId}/status", auth.agent_key)
        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, "uploadId": uploadId, "durationMs": round(duration_ms)})
            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)})

        output = {
            "uploadId": data.get("uploadId"),
            "status": data.get("status"),
            "ready": data.get("ready", False),
        }
        return json.dumps(output)
