import os
from typing import List, Callable, Any, Union
from functools import wraps
from dataclasses import dataclass
import httpx


@dataclass
class Value:
    """Successful function execution result."""
    data: Any


@dataclass
class Rejected:
    """Rejected function call with reason."""
    reason: str


Maybe = Union[Value, Rejected]


def guard(
    rules: List[str],
    model: str = "claude-haiku-4-5-20251001",
    max_tokens: int = 1024,
    use_thinking: bool = False,
    thinking_tokens: int = 2000
):
    """
    Decorator that validates function calls against rules using a remote Guard server.

    Args:
        rules: List of rules that the function call must comply with
        model: Claude model to use for validation
        max_tokens: Maximum tokens for the validation response
        use_thinking: Enable extended thinking for complex validations
        thinking_tokens: Budget for thinking tokens when use_thinking is True
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs) -> Maybe:
            # Get server configuration from environment
            server_url = os.environ.get("MORES_SERVER_URL", "http://localhost:8000")
            api_key = os.environ.get("MORES_API_KEY")

            if not api_key:
                raise RuntimeError("MORES_API_KEY environment variable not set")

            # Get function metadata
            func_name = func.__name__
            func_doc = func.__doc__ or "No documentation available"

            # Prepare the validation request
            request_data = {
                "function_name": func_name,
                "docstring": func_doc,
                "args": str(args),
                "kwargs": str(kwargs),
                "rules": rules,
                "model": model,
                "max_tokens": max_tokens,
                "use_thinking": use_thinking,
                "thinking_tokens": thinking_tokens
            }

            headers = {
                "X-API-Key": api_key,
                "Content-Type": "application/json"
            }

            try:
                # Make request to validation server
                response = httpx.post(
                    f"{server_url}/validate",
                    json=request_data,
                    headers=headers,
                    timeout=30.0
                )

                # Handle authentication errors
                if response.status_code == 401:
                    raise RuntimeError("Invalid API key")

                # Handle server errors
                if response.status_code >= 400:
                    error_detail = response.json().get("detail", "Unknown error")
                    raise RuntimeError(f"Server error: {error_detail}")

                # Parse validation response
                validation = response.json()

                # Check if the call is allowed
                if not validation.get("allowed", False):
                    return Rejected(reason=validation.get("reason", "No reason provided"))

                # If allowed, execute the function
                result = func(*args, **kwargs)
                return Value(data=result)

            except httpx.ConnectError:
                raise RuntimeError(
                    f"Cannot connect to Guard server at {server_url}. "
                    "Is the server running?"
                )
            except httpx.TimeoutException:
                raise RuntimeError("Request to Guard server timed out")
            except RuntimeError:
                raise  # Re-raise RuntimeError as-is
            except Exception as e:
                raise RuntimeError(f"Guard validation failed: {str(e)}")

        return wrapper
    return decorator
