from __future__ import annotations

import asyncio
import json
import time
from typing import Annotated, Any

from mcp.server.fastmcp import FastMCP

from ..api_client import ApiResult, api_get, api_post_json, api_post_multipart
from ..auth import resolve_auth
from ..config import get_transport_mode, mcp_config
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 ..security.wait_gate import acquire_wait_slot
from ..validation.base64_content import validate_and_decode_base64
from ..validation.file_path import validate_and_read_file
from ..validation.url import pre_validate_url

TOOL = "upload_and_wait"
_TERMINAL_STATUSES = frozenset({"READY", "FAILED", "DELETED"})
_BACKEND_DOWN_STATUSES = frozenset({502, 503, 504})
_MAX_CONSECUTIVE_ERRORS = 3


async def _poll_until_ready(
    upload_id: str,
    agent_key: str,
    deadline_ms: float,
) -> dict[str, Any]:
    start = time.monotonic()
    interval_ms = mcp_config.wait.poll_initial_ms
    polls = 0
    consecutive_errors = 0

    while (time.monotonic() - start) * 1000 < deadline_ms:
        await asyncio.sleep(interval_ms / 1000)
        polls += 1

        result = await api_get(f"/v1/upload/{upload_id}/status", agent_key)

        if result.ok:
            consecutive_errors = 0
            data = result.data
            if data.get("ready") or data.get("status") in _TERMINAL_STATUSES:
                return {"status": data["status"], "ready": data.get("ready", False), "polls": polls}
        else:
            consecutive_errors += 1
            if result.status in _BACKEND_DOWN_STATUSES:
                return {
                    "status": "SERVICE_UNAVAILABLE",
                    "ready": False,
                    "polls": polls,
                    "retryAfterSec": getattr(result, "retry_after_sec", None) or 30,
                }
            if consecutive_errors >= _MAX_CONSECUTIVE_ERRORS:
                return {
                    "status": "POLL_ERROR",
                    "ready": False,
                    "polls": polls,
                    "retryAfterSec": getattr(result, "retry_after_sec", None),
                }

        interval_ms = min(interval_ms * mcp_config.wait.poll_backoff_factor, mcp_config.wait.poll_max_ms)

    return {"status": "TIMEOUT", "ready": False, "polls": polls}


def register(mcp: FastMCP) -> None:
    @mcp.tool(
        name="upload_and_wait",
        description=(
            "Best tool for AI agents. Upload a file and wait until it is fully processed and ready for download — "
            "returns a working download URL in a single call. No polling needed. "
            "Provide EXACTLY ONE of: 'filePath' (local stdio only — zero tokens), 'url' (preferred for remote), or 'content' (base64 fallback). "
            "The server handles the entire upload + processing pipeline and returns only when the file is downloadable or an error occurs."
        ),
    )
    async def upload_and_wait(
        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,
        url: Annotated[str | None, "Public HTTPS URL of the file to upload (preferred for remote — avoids sending file data as tool argument)"] = None,
        content: Annotated[str | None, "Base64-encoded file content (use only when no URL or filePath is available)"] = None,
        fileName: Annotated[str | None, "Filename with extension. Required when using 'content', optional for 'url'."] = 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, "url": url, "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)

        slot_ok, slot, slot_err = acquire_wait_slot()
        if not slot_ok:
            audit_log({"tool": TOOL, "tokenPrefix": auth.token_prefix, "result": "error", "errorCode": "concurrency_limit"})
            raise ValueError(slot_err)

        try:
            return await _execute(args, auth.agent_key, auth.token_prefix, start, filePath, url, content, fileName, mimeType, pricePerDownload, idempotencyKey)
        finally:
            if slot:
                slot.release()


async def _execute(
    args: dict[str, Any],
    agent_key: str,
    token_prefix: str,
    start: float,
    filePath: str | None,
    url: str | None,
    content: str | None,
    fileName: str | None,
    mimeType: str | None,
    pricePerDownload: float | None,
    idempotencyKey: str | None,
) -> str:
    has_file_path = bool(filePath)
    has_url = bool(url)
    has_content = bool(content)

    source_count = sum([has_file_path, has_url, has_content])
    if source_count == 0:
        raise ValueError("Provide exactly one of 'filePath', 'url', or 'content'. None was supplied.")
    if source_count > 1:
        raise ValueError("Provide exactly one of 'filePath', 'url', or 'content'. Multiple were supplied.")

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

    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)

    upload_result: ApiResult

    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": token_prefix, "result": "error", "errorCode": "validation", "argsDigest": sanitize_args(TOOL, args)})
            raise ValueError(err)
        upload_result = await api_post_multipart("/v1/agent/upload", agent_key, data, name, mimeType or "application/octet-stream", extra_headers)  # type: ignore[arg-type]
    elif has_url:
        url_ok, url_err = pre_validate_url(url)  # type: ignore[arg-type]
        if not url_ok:
            audit_log({"tool": TOOL, "tokenPrefix": 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
        upload_result = await api_post_json("/v1/agent/upload-url", agent_key, body, extra_headers)  # type: ignore[arg-type]
    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": token_prefix, "result": "error", "errorCode": "validation", "argsDigest": sanitize_args(TOOL, args)})
            raise ValueError(err)
        upload_result = await api_post_multipart("/v1/agent/upload", agent_key, decoded, fileName, mimeType or "application/octet-stream", extra_headers)  # type: ignore[arg-type]

    if not upload_result.ok:
        duration_ms = (time.monotonic() - start) * 1000
        record_tool_call(TOOL, duration_ms, True)
        retry_hint = f" Retry after {upload_result.retry_after_sec}s." if upload_result.retry_after_sec else ""
        audit_log({"tool": TOOL, "tokenPrefix": token_prefix, "result": "error", "errorCode": upload_result.code, "durationMs": round(duration_ms), "argsDigest": sanitize_args(TOOL, args)})
        raise ValueError(upload_result.message + retry_hint)

    data = upload_result.data
    upload_id = data.get("uploadId", "")
    final_status = data.get("status", "PENDING_UPLOAD")
    ready = final_status == "READY"
    polls = 0
    retry_after_sec: int | None = None

    if not ready and final_status not in _TERMINAL_STATUSES:
        elapsed_ms = (time.monotonic() - start) * 1000
        remaining_ms = mcp_config.wait.max_duration_ms - elapsed_ms
        if remaining_ms > mcp_config.wait.poll_initial_ms:
            poll = await _poll_until_ready(upload_id, agent_key, remaining_ms)
            final_status = poll["status"]
            ready = poll["ready"]
            polls = poll["polls"]
            retry_after_sec = poll.get("retryAfterSec")

    duration_ms = (time.monotonic() - start) * 1000
    record_tool_call(TOOL, duration_ms, not ready)
    audit_log({
        "tool": TOOL,
        "tokenPrefix": token_prefix,
        "result": "ok" if ready else "error",
        "uploadId": upload_id,
        "polls": polls,
        **({"errorCode": final_status.lower()} if not ready else {}),
        "durationMs": round(duration_ms),
        "argsDigest": sanitize_args(TOOL, args),
    })

    output = {
        "uploadId": upload_id,
        "status": final_status,
        "ready": ready,
        "downloadUrl": data.get("downloadUrl"),
        "expiresAt": data.get("expiresAt"),
    }
    clean_output = {k: v for k, v in output.items() if v is not None}

    if not ready and final_status in ("SERVICE_UNAVAILABLE", "POLL_ERROR"):
        suffix = f" Storage backend is degraded — retry in {retry_after_sec}s." if retry_after_sec else " Storage backend is temporarily unavailable."
        raise ValueError(json.dumps(clean_output) + "\n\n" + suffix)

    if not ready and final_status == "FAILED":
        raise ValueError(json.dumps(clean_output) + "\n\nUpload failed during processing.")

    result_str = json.dumps(clean_output)
    if not ready and final_status == "TIMEOUT":
        return result_str + "\n\nFile is still processing. Use get_upload_status to check later."

    return result_str
