# Copyright (c) 2024, CoSig Contributors
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""FastMCP integration for CoSig - YubiKey approval for MCP tools.

This module provides a decorator-based integration with FastMCP that requires
YubiKey approval before executing sensitive tools.

Usage:

    from fastmcp import FastMCP
    from cosig import configure, require_approval

    # Configure to use CoSig cloud (or self-hosted instance)
    configure(
        api_url="https://api.cosig.io",
        api_key="sk-your-api-key",
        username="your-username",
    )

    mcp = FastMCP("My Server")

    @mcp.tool()
    @require_approval()
    async def delete_customer(customer_id: str) -> str:
        '''Delete customer - requires YubiKey approval'''
        await db.delete(customer_id)
        return f"Deleted {customer_id}"

    # The first call will raise ApprovalRequiredError
    # User must tap YubiKey at the approval URL
    # Then retry the call - it will execute

Flow:
    1. Client calls tool
    2. Decorator checks for approval (hash of tool+args)
    3. If not approved: raises ApprovalRequiredError with URL
    4. User opens URL, taps YubiKey
    5. Client retries tool call
    6. Decorator sees approval, executes tool

Alternative Flow with wait_for_approval=True (recommended for AI agents):
    1. User keeps the CoSig dashboard open in browser (shows pending approvals)
    2. AI calls tool
    3. Tool starts polling immediately (appears "running" in AI chat)
    4. Pending approval appears on dashboard (auto-refreshes)
    5. User clicks approval, taps YubiKey
    6. Polling detects approval, executes tool, returns result
    7. No manual intervention needed - fully automatic!
"""

import asyncio
import functools
import hashlib
import sys
from collections.abc import Callable
from typing import Any

from cosig.security.canonical_json import canonicalize_json

# Approval expiry window - approvals older than this require fresh approval
APPROVAL_EXPIRY_MINUTES = 5


class ApprovalRequiredError(Exception):
    """Raised when a tool requires YubiKey approval before execution.

    The MCP client should display the approval_url to the user and retry
    the tool call after the user has tapped their YubiKey.
    """

    def __init__(self, message: str, plan_hash: str, approval_url: str):
        """Initialize approval required error.

        Args:
            message: Human-readable error message
            plan_hash: Hash of the execution plan
            approval_url: URL where user can approve with YubiKey
        """
        super().__init__(message)
        self.plan_hash = plan_hash
        self.approval_url = approval_url


def compute_tool_hash(tool_name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
    """Compute deterministic hash of a tool call.

    This must match the server's compute_plan_hash() function exactly.
    The format is: {"args": combined_arguments, "tool": tool_name}
    with compact JSON (no spaces) and sorted keys.

    Args:
        tool_name: Name of the tool function
        args: Positional arguments
        kwargs: Keyword arguments

    Returns:
        SHA256 hash of the tool call
    """
    # Combine args and kwargs into a single arguments dict
    # Positional args are stored with indexed keys (matches cloud_client.py)
    arguments = dict(kwargs)
    for i, arg in enumerate(args):
        arguments[f"_arg{i}"] = arg

    # Create canonical representation matching server format
    call_data = {
        "tool": tool_name,
        "args": arguments,
    }

    # Use canonical JSON for deterministic hashing
    canonical = canonicalize_json(call_data)
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def require_approval(
    username: str | None = None,
    wait_for_approval: bool = False,
    poll_interval: float = 2.0,
    timeout: float = 300.0,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Decorator that requires YubiKey approval before executing a tool.

    This decorator integrates with FastMCP tools to require hardware token
    approval before execution. It communicates with the CoSig cloud service
    (or self-hosted instance) to manage approval requests.

    Two modes of operation:

    1. wait_for_approval=False (default): Raises ApprovalRequiredError immediately.
       User must approve and manually retry the tool call.

    2. wait_for_approval=True: Polls for approval automatically.
       User approves via YubiKey and the tool executes without manual retry.
       This provides a smoother UX for AI agent interactions.

    Args:
        username: Email address of the user who should approve this request.
                 If not provided, will look for 'username' in tool kwargs, or default to
                 the API key's associated user. For multi-tenant apps, pass the user's email.
        wait_for_approval: If True, poll for approval instead of raising error
        poll_interval: Seconds between polling checks (default: 2.0)
        timeout: Maximum seconds to wait for approval (default: 300.0 = 5 min)

    Returns:
        Decorator function

    Example (with username per decorator):
        @mcp.tool()
        @require_approval(username="alice")
        async def dangerous_operation(data: str) -> str:
            # This only executes after YubiKey approval by alice
            return perform_operation(data)

    Example (with username passed at runtime):
        @mcp.tool()
        @require_approval()
        async def dangerous_operation(username: str, data: str) -> str:
            # Username extracted from function args
            return perform_operation(data)

    Example (polling mode):
        @mcp.tool()
        @require_approval(wait_for_approval=True, timeout=120)
        async def dangerous_operation(username: str, data: str) -> str:
            # Waits up to 2 minutes for approval, then executes
            return perform_operation(data)
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        """Wrap function with approval check."""

        @functools.wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            """Check for approval before executing."""
            from cosig.cloud_client import CloudClientError, get_cloud_client
            from cosig.config import get_config

            # Extract FastMCP context if present (for progress notifications)
            ctx = kwargs.pop("ctx", None)

            # Determine username for this approval request
            # Priority: decorator param > function kwarg > None (uses API key's user)
            effective_username = username
            if effective_username is None and "username" in kwargs:
                effective_username = kwargs["username"]
            # If still None, the backend will use the API key's associated user

            config = get_config()
            client = get_cloud_client()

            # Compute hash of this tool call
            tool_name = func.__name__
            plan_hash = compute_tool_hash(tool_name, args, kwargs)

            # Build approval URL from config
            # The approval URL is on the web frontend, not the API
            base_url = config.api_url.rstrip("/")

            # Convert API URL to web URL
            # Example: https://api.cosig.io/api/v1 -> https://cosig.io
            web_url = base_url

            # First, convert api.domain.com to domain.com
            if "://api." in web_url:
                web_url = web_url.replace("://api.", "://")

            # Then remove API path suffixes
            if web_url.endswith("/api/v1"):
                web_url = web_url[:-7]
            elif web_url.endswith("/api"):
                web_url = web_url[:-4]

            approval_url = f"{web_url}/approve/{plan_hash}"

            try:
                # Check approval status via cloud API
                status_response = await client.check_approval_status(plan_hash)
                approval_status = status_response.get("status", "not_found")
            except CloudClientError as e:
                if e.status_code == 404:
                    approval_status = "not_found"
                else:
                    raise

            # Create a new approval request if:
            # - No plan exists (not_found)
            # - Previous plan was rejected (user can try again with fresh request)
            # - Previous plan was executed (new execution needs fresh approval)
            # Only "pending" plans should prevent creating a new request
            if approval_status in ("not_found", "rejected", "executed"):
                # Create new approval request via cloud API
                await client.create_approval_request(
                    plan_hash=plan_hash,
                    tool_name=tool_name,
                    args=list(args),
                    kwargs=kwargs,
                    username=effective_username,
                )

                if wait_for_approval:
                    # Start polling immediately
                    return await _wait_and_execute(
                        func=func,
                        args=args,
                        kwargs=kwargs,
                        plan_hash=plan_hash,
                        poll_interval=poll_interval,
                        timeout=timeout,
                        approval_url=approval_url,
                        tool_name=tool_name,
                        ctx=ctx,
                        client=client,
                    )
                else:
                    # Traditional mode: raise error for manual retry
                    raise ApprovalRequiredError(
                        message=(
                            f"YubiKey approval required for {tool_name}. "
                            f"Tap your YubiKey at {approval_url} then retry."
                        ),
                        plan_hash=plan_hash,
                        approval_url=approval_url,
                    )

            elif approval_status == "pending":
                if wait_for_approval:
                    # Continue polling
                    return await _wait_and_execute(
                        func=func,
                        args=args,
                        kwargs=kwargs,
                        plan_hash=plan_hash,
                        poll_interval=poll_interval,
                        timeout=timeout,
                        approval_url=approval_url,
                        tool_name=tool_name,
                        ctx=ctx,
                        client=client,
                    )
                else:
                    # Traditional mode: raise error for manual retry
                    raise ApprovalRequiredError(
                        message=(
                            f"YubiKey approval pending for {tool_name}. "
                            f"Tap your YubiKey at {approval_url} then retry."
                        ),
                        plan_hash=plan_hash,
                        approval_url=approval_url,
                    )

            elif approval_status == "approved":
                # Approved - execute the tool and mark as executed
                result = await func(*args, **kwargs)
                await client.mark_executed(plan_hash, str(result))
                return result

            else:
                raise ValueError(f"Unknown approval status: {approval_status}")

        return wrapper

    return decorator


async def _wait_and_execute(
    func: Callable[..., Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    plan_hash: str,
    poll_interval: float,
    timeout: float,
    approval_url: str,
    tool_name: str,
    ctx: Any = None,
    client: Any = None,
) -> Any:
    """Poll for approval and execute the tool when approved.

    Args:
        func: The tool function to execute
        args: Positional arguments for the tool
        kwargs: Keyword arguments for the tool
        plan_hash: Hash of the tool call
        poll_interval: Seconds between polling checks
        timeout: Maximum seconds to wait for approval
        approval_url: URL for approval (for error messages)
        tool_name: Name of the tool (for error messages)
        ctx: FastMCP context for progress notifications (optional)
        client: Cloud client instance

    Returns:
        Result from the tool execution

    Raises:
        ApprovalRequiredError: If timeout is reached without approval
    """
    elapsed = 0.0
    print(f"\n[COSIG] YubiKey approval required for {tool_name}", file=sys.stderr)
    print(f"[COSIG] Please approve at: {approval_url}", file=sys.stderr)
    print(f"[COSIG] Waiting for approval... (timeout: {timeout}s)\n", file=sys.stderr)

    # Send progress notification if context is available
    if ctx:
        try:
            if hasattr(ctx, "info"):
                await ctx.info(f"YubiKey approval required for {tool_name}")
                await ctx.info(f"Please visit: {approval_url}")
            if hasattr(ctx, "report_progress"):
                await ctx.report_progress(0, 100, f"Waiting for YubiKey approval at {approval_url}")
        except Exception as e:
            print(f"[COSIG] Could not send progress notification: {e}", file=sys.stderr)

    while elapsed < timeout:
        await asyncio.sleep(poll_interval)
        elapsed += poll_interval

        try:
            status_response = await client.check_approval_status(plan_hash)
            status = status_response.get("status", "pending")
        except Exception:
            status = "pending"

        # Update progress periodically
        if ctx and hasattr(ctx, "report_progress"):
            try:
                progress = min(int((elapsed / timeout) * 100), 99)
                await ctx.report_progress(
                    progress, 100, f"Waiting for YubiKey approval... ({int(elapsed)}s)"
                )
            except Exception:
                pass

        if status == "approved":
            print(f"\n[COSIG] Approval received! Executing {tool_name}...\n", file=sys.stderr)
            result = await func(*args, **kwargs)
            await client.mark_executed(plan_hash, str(result))
            return result

        elif status == "executed":
            # Already executed by another process
            print(f"\n[COSIG] Already executed, re-running {tool_name}...\n", file=sys.stderr)
            return await func(*args, **kwargs)

        elif status == "rejected":
            # User rejected the approval request
            print(f"\n[COSIG] Approval REJECTED for {tool_name}.\n", file=sys.stderr)
            return (
                f"Tool execution rejected by user. The approval request for {tool_name} was denied."
            )

        elif status == "not_found":
            # Plan was deleted? This shouldn't happen
            raise ApprovalRequiredError(
                message="Approval request was deleted. Please retry the tool call.",
                plan_hash=plan_hash,
                approval_url=approval_url,
            )

        # status == "pending" - continue polling

    # Timeout reached - return a helpful message
    return (
        f"Approval timeout for {tool_name}.\n\n"
        f"Please make sure you have the CoSig dashboard open in your browser.\n"
        f"Pending approvals appear there automatically.\n\n"
        f"Once you've approved, run this command again."
    )
