import os
import inspect
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


@dataclass
class DoubtError(Exception):
    """Raised when doubt validation fails."""

    reason: str


Maybe = Union[Value, Rejected]


def guard(rules: List[str], mode: str = "fast", raise_reject: bool = False):
    """
    Decorator that validates function calls against rules using a remote server.

    Args:
        rules: List of rules that the function call must comply with
        mode: "fast" (Haiku 4.5, no thinking) or "slow" (Sonnet 4.5, with thinking)
        raise_reject: If True, raises RuntimeError on rejection instead of returning Rejected
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Get server configuration from environment
            server_url = os.environ.get("MORES_SERVER_URL", "https://mores.fulcrumresearch.ai")
            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,
                "mode": mode,
            }

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

            try:
                # Make request to guard server
                response = httpx.post(
                    f"{server_url}/guard",
                    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):
                    reason = validation.get("reason", "No reason provided")
                    if raise_reject:
                        raise RuntimeError(f"Guard rejected: {reason}")
                    return Rejected(reason=reason)

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

            except httpx.ConnectError:
                raise RuntimeError(
                    f"Cannot connect to Mores server at {server_url}. "
                    "Is the server running?"
                )
            except httpx.TimeoutException:
                raise RuntimeError("Request to Mores 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


def mock(mode: str = "fast"):
    """
    Decorator that replaces function execution with LLM-generated mock output.

    Args:
        mode: "fast" (Haiku 4.5, no thinking) or "slow" (Sonnet 4.5, with thinking)
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Get server configuration from environment
            server_url = os.environ.get("MORES_SERVER_URL", "https://mores.fulcrumresearch.ai")
            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"
            source_code = inspect.getsource(func)

            # Prepare the mock request
            request_data = {
                "function_name": func_name,
                "docstring": func_doc,
                "args": str(args),
                "kwargs": str(kwargs),
                "source_code": source_code,
                "mode": mode,
            }

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

            try:
                # Make request to mock server
                response = httpx.post(
                    f"{server_url}/mock",
                    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 mock response
                mock_result = response.json()
                return mock_result.get("result", "")

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

        return wrapper

    return decorator


def doubt(mode: str = "fast"):
    """
    Decorator that validates function output correctness post-execution using LLM.

    Raises DoubtError if the output is determined to be incorrect.

    Args:
        mode: "fast" (Haiku 4.5, no thinking) or "slow" (Sonnet 4.5, with thinking)
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Execute function first
            result = func(*args, **kwargs)

            # Get server configuration from environment
            server_url = os.environ.get("MORES_SERVER_URL", "https://mores.fulcrumresearch.ai")
            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"
            source_code = inspect.getsource(func)

            # Prepare the doubt request
            request_data = {
                "function_name": func_name,
                "docstring": func_doc,
                "args": str(args),
                "kwargs": str(kwargs),
                "source_code": source_code,
                "result": str(result),
                "mode": mode,
            }

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

            try:
                # Make request to doubt server
                response = httpx.post(
                    f"{server_url}/doubt",
                    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 doubt response
                validation = response.json()

                # If not valid, raise DoubtError
                if not validation.get("allowed", False):
                    raise DoubtError(
                        validation.get("reason", "Output validation failed")
                    )

                # If valid, return result
                return result

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

        return wrapper

    return decorator
