"""Core data models and interfaces for the MCP Proxy.

This module defines the data structures used throughout the proxy system,
including JSON-RPC messages, HTTP requests/responses, authentication state,
and configuration models.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field, HttpUrl, field_validator


class MessageStatus(str, Enum):
    """Status of message processing."""
    PENDING = "pending"
    COMPLETED = "completed"
    FAILED = "failed"


# ============================================================================
# JSON-RPC Models
# ============================================================================


class JsonRpcError(BaseModel):
    """JSON-RPC 2.0 error object."""
    code: int = Field(..., description="Error code")
    message: str = Field(..., description="Error message")
    data: Optional[Any] = Field(None, description="Additional error data")


class JsonRpcMessage(BaseModel):
    """JSON-RPC 2.0 message structure."""
    jsonrpc: str = Field(default="2.0", description="JSON-RPC version")
    id: Optional[Union[str, int]] = Field(None, description="Request/response ID")
    method: Optional[str] = Field(None, description="Method name for requests")
    params: Optional[Any] = Field(None, description="Method parameters")
    result: Optional[Any] = Field(None, description="Result for responses")
    error: Optional[JsonRpcError] = Field(None, description="Error for error responses")

    @field_validator("jsonrpc")
    @classmethod
    def validate_jsonrpc_version(cls, v: str) -> str:
        """Ensure JSON-RPC version is 2.0."""
        if v != "2.0":
            raise ValueError("JSON-RPC version must be 2.0")
        return v

    def is_request(self) -> bool:
        """Check if this is a request message."""
        return self.method is not None

    def is_response(self) -> bool:
        """Check if this is a response message."""
        return self.result is not None or self.error is not None

    def is_notification(self) -> bool:
        """Check if this is a notification (request without ID)."""
        return self.method is not None and self.id is None

    def to_json(self) -> str:
        """Serialize to JSON string."""
        return self.model_dump_json(exclude_none=True)

    @classmethod
    def from_json(cls, json_str: str) -> JsonRpcMessage:
        """Parse from JSON string."""
        return cls.model_validate_json(json_str)


# ============================================================================
# HTTP MCP Models
# ============================================================================


class HttpMcpRequest(BaseModel):
    """HTTP MCP request structure."""
    method: str = Field(..., description="HTTP method (POST, GET)")
    url: str = Field(..., description="Request URL")
    headers: Dict[str, str] = Field(default_factory=dict, description="HTTP headers")
    body: Optional[str] = Field(None, description="Request body")
    session_id: Optional[str] = Field(None, description="MCP session ID")

    @field_validator("method")
    @classmethod
    def validate_method(cls, v: str) -> str:
        """Ensure HTTP method is valid."""
        if v.upper() not in ["GET", "POST", "PUT", "DELETE", "PATCH"]:
            raise ValueError(f"Invalid HTTP method: {v}")
        return v.upper()


class HttpMcpResponse(BaseModel):
    """HTTP MCP response structure."""
    status: int = Field(..., description="HTTP status code")
    headers: Dict[str, str] = Field(default_factory=dict, description="Response headers")
    body: str = Field(default="", description="Response body")
    content_type: str = Field(default="application/json", description="Content type")

    def is_success(self) -> bool:
        """Check if response indicates success."""
        return 200 <= self.status < 300

    def is_auth_error(self) -> bool:
        """Check if response indicates authentication error."""
        return self.status == 401


# ============================================================================
# Authentication Models
# ============================================================================


class ClientCredentials(BaseModel):
    """OAuth client credentials."""
    client_id: str = Field(..., description="OAuth client ID")
    client_secret: str = Field(..., description="OAuth client secret")
    expires_at: Optional[datetime] = Field(None, description="Credential expiration time")

    def is_expired(self) -> bool:
        """Check if credentials are expired."""
        if self.expires_at is None:
            return False
        return datetime.now() >= self.expires_at


class OAuthTokenResponse(BaseModel):
    """OAuth token response."""
    access_token: str = Field(..., description="Access token")
    token_type: str = Field(..., description="Token type (usually Bearer)")
    expires_in: int = Field(..., description="Token lifetime in seconds")
    scope: Optional[str] = Field(None, description="Token scope")
    refresh_token: Optional[str] = Field(None, description="Refresh token")

    def get_expiration_time(self) -> datetime:
        """Calculate token expiration time."""
        from datetime import timedelta
        return datetime.now() + timedelta(seconds=self.expires_in)


class AuthenticationState(BaseModel):
    """Current authentication state."""
    client_credentials: Optional[ClientCredentials] = Field(
        None, description="Client credentials"
    )
    access_token: Optional[str] = Field(None, description="Current access token")
    refresh_token: Optional[str] = Field(None, description="Refresh token for token renewal")
    token_expires_at: Optional[datetime] = Field(
        None, description="Token expiration time"
    )
    last_dcr_attempt: Optional[datetime] = Field(
        None, description="Last DCR attempt timestamp"
    )
    dcr_retry_count: int = Field(default=0, description="DCR retry counter")
    is_authenticated: bool = Field(default=False, description="Authentication status")

    def is_token_valid(self) -> bool:
        """Check if current token is valid."""
        if not self.access_token or not self.token_expires_at:
            return False
        # Consider token invalid 60 seconds before actual expiration
        from datetime import timedelta
        return datetime.now() < (self.token_expires_at - timedelta(seconds=60))

    def needs_token_refresh(self) -> bool:
        """Check if token needs refresh."""
        return self.is_authenticated and not self.is_token_valid()


# ============================================================================
# Session Models
# ============================================================================


class SessionState(BaseModel):
    """MCP session state."""
    session_id: Optional[str] = Field(None, description="Session ID")
    is_initialized: bool = Field(default=False, description="Initialization status")
    protocol_version: str = Field(default="1.0", description="MCP protocol version")
    capabilities: List[str] = Field(default_factory=list, description="Session capabilities")
    sse_stream_active: bool = Field(default=False, description="SSE stream status")


# ============================================================================
# Message Correlation Models
# ============================================================================


@dataclass
class MessageCorrelation:
    """Tracks correlation between stdio and HTTP messages."""
    stdio_request_id: Union[str, int]
    http_request_id: str
    timestamp: datetime
    method: str
    status: MessageStatus = MessageStatus.PENDING

    def is_pending(self) -> bool:
        """Check if message is still pending."""
        return self.status == MessageStatus.PENDING

    def mark_completed(self) -> None:
        """Mark message as completed."""
        self.status = MessageStatus.COMPLETED

    def mark_failed(self) -> None:
        """Mark message as failed."""
        self.status = MessageStatus.FAILED


# ============================================================================
# Configuration Models
# ============================================================================


class ProxyConfig(BaseModel):
    """Proxy configuration."""
    mcp_server_url: HttpUrl = Field(..., description="HTTP MCP server URL")
    oauth_provider_url: HttpUrl = Field(..., description="OAuth provider URL")
    client_name: str = Field(default="mcp-proxy-client", description="OAuth client name")
    scopes: List[str] = Field(
        default_factory=lambda: ["mcp:read", "mcp:write"],
        description="OAuth scopes"
    )
    connection_timeout: int = Field(
        default=30, ge=1, le=300, description="Connection timeout in seconds"
    )
    retry_attempts: int = Field(
        default=3, ge=0, le=10, description="Number of retry attempts"
    )
    log_level: str = Field(
        default="info", description="Logging level"
    )
    max_backoff_seconds: int = Field(
        default=60, ge=1, le=300, description="Maximum backoff time in seconds"
    )

    @field_validator("log_level")
    @classmethod
    def validate_log_level(cls, v: str) -> str:
        """Ensure log level is valid."""
        valid_levels = ["debug", "info", "warning", "error", "critical"]
        if v.lower() not in valid_levels:
            raise ValueError(f"Invalid log level: {v}. Must be one of {valid_levels}")
        return v.lower()

    @field_validator("scopes")
    @classmethod
    def validate_scopes(cls, v: List[str]) -> List[str]:
        """Ensure scopes list is not empty."""
        if not v:
            raise ValueError("Scopes list cannot be empty")
        return v