"""
JSON-RPC Protocol Definitions

Message schema for communication between AWS backend and workspace agent.
All messages follow JSON-RPC 2.0 over WebSocket.
"""
from __future__ import annotations

from dataclasses import dataclass, field, asdict
from typing import Optional, Dict, Any, List
from enum import Enum
import json
import uuid


class MessageType(str, Enum):
    """Control message types (non-RPC)."""
    HEARTBEAT = "heartbeat"
    HEARTBEAT_ACK = "heartbeat_ack"
    AUTH = "auth"
    AUTH_OK = "auth_ok"
    AUTH_FAIL = "auth_fail"
    CREDENTIAL_SYNC = "credential_sync"
    CREDENTIAL_ACK = "credential_ack"
    CREDENTIAL_DELETE = "credential_delete"
    POLICY_UPDATE = "policy_update"
    UPDATE_AVAILABLE = "update_available"
    DISCONNECT = "disconnect"
    # Signal bridge
    START_SIGNAL_BRIDGE = "start_signal_bridge"
    STOP_SIGNAL_BRIDGE = "stop_signal_bridge"
    SYNC_SIGNAL_BRIDGES = "sync_signal_bridges"
    SIGNAL_THROTTLE = "signal_throttle"


class ToolType(str, Enum):
    """Tool execution modes."""
    LOCAL = "local"       # filesystem, shell, browser, etc.
    SERVICE = "service"   # HTTP proxy with credential resolution


@dataclass
class AuthInfo:
    """Authentication descriptor for service execution specs.
    Tells the workspace agent HOW to inject credentials — not the credentials themselves."""
    service_slug: str
    auth_type: str          # "bearer", "apikey", "basic", "custom"
    header_name: str = "Authorization"
    header_value_pattern: str = "{value}"  # e.g. "Bearer {value}"
    query_param_name: Optional[str] = None
    username: Optional[str] = None  # For basic auth


@dataclass
class ExecutionSpec:
    """Fully-built HTTP request spec — everything except credentials."""
    url: str
    method: str = "POST"
    headers: Dict[str, str] = field(default_factory=dict)
    body: Optional[str] = None
    timeout: int = 30
    max_response_bytes: int = 1_048_576  # 1MB
    auth: Optional[AuthInfo] = None


@dataclass
class AgentLimits:
    """Per-agent execution limits. Sent from backend, enforced on workspace."""
    agent_id: str = ""
    agent_name: str = ""
    enabled_tools: List[str] = field(default_factory=list)
    allowed_paths: List[str] = field(default_factory=lambda: ["/workspace"])
    blocked_paths: List[str] = field(default_factory=lambda: ["**/.env", "**/.vault"])
    allowed_commands: List[str] = field(default_factory=list)  # empty = all allowed
    blocked_commands: List[str] = field(default_factory=lambda: ["rm -rf /", "sudo", "shutdown"])
    max_shell_timeout: int = 300
    max_concurrent_shells: int = 3
    max_requests_per_minute: int = 60
    allowed_services: List[str] = field(default_factory=list)  # empty = all allowed
    max_browser_sessions: int = 2
    max_file_size_bytes: int = 52_428_800  # 50MB


@dataclass
class ToolRequest:
    """Inbound tool execution request from backend."""
    id: str
    tool: str
    tool_type: ToolType = ToolType.LOCAL
    agent_id: str = ""
    agent_name: str = ""
    arguments: Dict[str, Any] = field(default_factory=dict)
    execution_spec: Optional[ExecutionSpec] = None
    limits: Optional[AgentLimits] = None
    timeout: int = 300

    @classmethod
    def from_jsonrpc(cls, msg: dict) -> "ToolRequest":
        """Parse a JSON-RPC execute_tool request."""
        params = msg.get("params", {})
        spec_data = params.get("execution_spec")
        spec = None
        if spec_data:
            auth_data = spec_data.pop("auth", None)
            auth = AuthInfo(**auth_data) if auth_data else None
            spec = ExecutionSpec(**spec_data, auth=auth)

        limits_data = params.get("limits")
        limits = AgentLimits(**limits_data) if limits_data else None

        return cls(
            id=msg["id"],
            tool=params["tool"],
            tool_type=ToolType(params.get("tool_type", "local")),
            agent_id=params.get("agent_id", ""),
            agent_name=params.get("agent_name", ""),
            arguments=params.get("arguments", {}),
            execution_spec=spec,
            limits=limits,
            timeout=params.get("timeout", 300),
        )


@dataclass
class ToolResponse:
    """Outbound tool execution result."""
    id: str
    success: bool
    output: str = ""
    error: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_jsonrpc(self) -> str:
        """Serialize to JSON-RPC response."""
        payload = {
            "jsonrpc": "2.0",
            "id": self.id,
            "result": {
                "success": self.success,
                "output": self.output,
                "error": self.error,
                "metadata": self.metadata,
            }
        }
        return json.dumps(payload)


@dataclass
class ControlMessage:
    """Non-RPC control messages (heartbeat, auth, credential sync, etc.)."""
    type: MessageType
    data: Dict[str, Any] = field(default_factory=dict)

    def to_json(self) -> str:
        return json.dumps({"type": self.type.value, **self.data})

    @classmethod
    def from_json(cls, raw: str) -> "ControlMessage":
        msg = json.loads(raw)
        msg_type = MessageType(msg.pop("type"))
        return cls(type=msg_type, data=msg)


def parse_message(raw: str) -> ToolRequest | ControlMessage:
    """Parse an incoming WSS message into the appropriate type."""
    msg = json.loads(raw)
    if "jsonrpc" in msg and msg.get("method") == "execute_tool":
        return ToolRequest.from_jsonrpc(msg)
    elif "type" in msg:
        return ControlMessage.from_json(raw)
    else:
        raise ValueError(f"Unknown message format: {list(msg.keys())}")


def make_auth_message(token: str, version: str) -> str:
    """Build the initial auth handshake message."""
    return json.dumps({
        "type": MessageType.AUTH.value,
        "token": token,
        "agent_version": version,
    })


def make_heartbeat() -> str:
    return json.dumps({"type": MessageType.HEARTBEAT.value})


def make_heartbeat_ack() -> str:
    return json.dumps({"type": MessageType.HEARTBEAT_ACK.value})
