import base64
import gzip
import hashlib
import json
from collections.abc import Awaitable, Callable
from enum import Enum
from typing import Literal, Protocol, TypeAlias
from uuid import uuid4

from pydantic import BaseModel, Field, JsonValue


class ExecuteCodeRequest(BaseModel):
    tool_names: list[str]
    code: str


class MessageType(str, Enum):
    TOOL_CALL = "tool_call"
    SUCCESS = "success"
    SYNTAX_ERROR = "syntax_error"
    RUNTIME_ERROR = "runtime_error"
    TIMEOUT_ERROR = "timeout_error"
    DETERMINISM_ERROR = "determinism_error"


def _tool_call_hash(
    tool_name: str,
    args: tuple[JsonValue, ...],
    kwargs: dict[str, JsonValue],
) -> str:
    data = json.dumps(
        {"tool_name": tool_name, "args": args, "kwargs": kwargs},
        sort_keys=True,
    )
    return hashlib.sha256(data.encode()).hexdigest()[:16]


class ToolCall(BaseModel):
    type: Literal[MessageType.TOOL_CALL] = MessageType.TOOL_CALL
    id: str = Field(default_factory=lambda: uuid4().hex)
    execution_id: str
    tool_name: str
    args: tuple[JsonValue, ...]
    kwargs: dict[str, JsonValue]

    def call_hash(self) -> str:
        return _tool_call_hash(self.tool_name, self.args, self.kwargs)


class ExecuteCodeSuccess(BaseModel):
    type: Literal[MessageType.SUCCESS] = MessageType.SUCCESS
    stdout: str
    stderr: str


class CodeSyntaxError(BaseModel):
    type: Literal[MessageType.SYNTAX_ERROR] = MessageType.SYNTAX_ERROR
    stdout: str
    stderr: str


class CodeRuntimeError(BaseModel):
    type: Literal[MessageType.RUNTIME_ERROR] = MessageType.RUNTIME_ERROR
    stdout: str
    stderr: str


class CodeTimeoutError(BaseModel):
    type: Literal[MessageType.TIMEOUT_ERROR] = MessageType.TIMEOUT_ERROR
    stdout: str
    stderr: str


class CodeDeterminismError(BaseModel):
    type: Literal[MessageType.DETERMINISM_ERROR] = MessageType.DETERMINISM_ERROR
    stdout: str
    stderr: str


ExecuteCodeError: TypeAlias = CodeSyntaxError | CodeRuntimeError | CodeTimeoutError


ExecuteCodeMessage: TypeAlias = ToolCall | ExecuteCodeSuccess | ExecuteCodeError


class ToolCallResult(BaseModel):
    execution_id: str
    tool_call_id: str
    result: JsonValue


ToolHandler: TypeAlias = Callable[..., Awaitable[JsonValue]]


class ToolCallReceivedEvent(BaseModel):
    type: Literal["tool_call_received"] = "tool_call_received"
    tool_call: ToolCall


class ToolCallCompletedEvent(BaseModel):
    type: Literal["tool_call_completed"] = "tool_call_completed"
    tool_call: ToolCall
    result: JsonValue


class ToolCallSentEvent(BaseModel):
    type: Literal["tool_call_sent"] = "tool_call_sent"
    tool_call: ToolCall
    result: JsonValue


ExecutionEvent: TypeAlias = (
    ToolCallReceivedEvent | ToolCallCompletedEvent | ToolCallSentEvent
)


class SyncExecutionEventHandler(Protocol):
    def __call__(self, event: ExecutionEvent) -> None: ...


class AsyncExecutionEventHandler(Protocol):
    def __call__(self, event: ExecutionEvent) -> Awaitable[None]: ...


ExecutionEventHandler: TypeAlias = (
    SyncExecutionEventHandler | AsyncExecutionEventHandler
)


class ExecuteCodeStepState(BaseModel):
    tool_results: list[JsonValue] = Field(default_factory=list)
    tool_call_hashes: list[str] = Field(default_factory=list)

    def encode(self) -> str:
        data = self.model_dump_json().encode()
        return base64.b64encode(gzip.compress(data)).decode()

    @classmethod
    def decode(cls, encoded: str) -> "ExecuteCodeStepState":
        data = gzip.decompress(base64.b64decode(encoded))
        return cls.model_validate_json(data)


class ExecuteCodeStepRequest(BaseModel):
    tool_names: list[str]
    code: str
    tool_results: list[JsonValue] | None = None
    state: str | None = None


class ExecuteCodeStepToolCalls(BaseModel):
    type: Literal[MessageType.TOOL_CALL] = MessageType.TOOL_CALL
    tool_calls: list[ToolCall]
    state: str


ExecuteCodeStepMessage: TypeAlias = (
    ExecuteCodeStepToolCalls
    | ExecuteCodeSuccess
    | ExecuteCodeError
    | CodeDeterminismError
)
