import uuid
from dataclasses import dataclass
from typing import Any, Literal, overload

from msgspec import Struct, field, to_builtins
from msgspec.json import Encoder

from ..core.datetime import time_now_utc

__all__ = [
    'CreateAgentRequest',
    'PromptTemplate',
    'should_template',
    'InteractionCodeBlock',
    'InteractionExecuteResult',
    'InteractionEvent',
    'ReasoningEffort',
    # sm messages
    'SMCreateAgentMessage',
    'SMDestroyAgentMessage',
    'SMInvocationEnterMessage',
    'SMInvocationExitMessage',
    'SMInvocationErrorMessage',
    'SessionManagerMessage',
    'SMInferenceRequestMessage',
    'SMInferenceResponseMessage',
    'SMInferenceErrorMessageMessage',
    'SMInvocationInteractionMessage',
    'SMMonadMessage',
    # monad messages
    'MonadStreamMessage',
    # magic messages
    'MagicMessage',
    'MagicExecuteEnterMessage',
    'MagicExecuteExitMessage',
    'MagicInferenceEnterMessage',
    'MagicInferenceExitMessage',
    'AllServerMessage',
    # encode
    'server_message_to_json',
    'server_message_to_dict',
]


# Request from the client

DEFAULT_PROTOCOL = 'python/v0'

# as per https://openrouter.ai/docs/guides/best-practices/reasoning-tokens#reasoning-effort-level
type ReasoningEffort = Literal['none', 'minimal', 'low', 'medium', 'high', 'xhigh']


class InteractionCodeBlock(Struct):
    code: str
    exec_id: str = field(default_factory=lambda: str(uuid.uuid4()))


class InteractionExecuteResult(Struct):
    result: str
    exec_id: str


type InteractionEvent = InteractionCodeBlock | InteractionExecuteResult


@dataclass(kw_only=True)
class PromptTemplate:
    template: str

    def __bool__(self) -> bool:
        return bool(self.template)


@overload
def should_template(value: PromptTemplate) -> str: ...
@overload
def should_template(value: str) -> None: ...
@overload
def should_template(value: None) -> None: ...
def should_template(value: str | PromptTemplate | None) -> str | None:
    if isinstance(value, str):
        return None
    if value is None:
        return None
    return value.template


@dataclass(kw_only=True)
class CreateAgentRequest:
    doc: str | None
    system: str | PromptTemplate | None
    model: str
    json: bool
    streaming: bool
    warp_globals_payload: bytes
    protocol: str = DEFAULT_PROTOCOL
    max_tokens_per_invocation: int | None = None
    max_tokens_per_round: int | None = None
    max_rounds: int | None = None
    reasoning_effort: ReasoningEffort | None = None


# Messages generated by the server


class SMCreateAgentMessage(Struct, tag="sm_create_agent"):
    uid: str
    doc: str | None
    system: str | None
    model: str
    json: bool
    streaming: bool
    timestamp: str = field(default_factory=time_now_utc)


class SMDestroyAgentMessage(Struct, tag="sm_destroy_agent"):
    uid: str
    timestamp: str = field(default_factory=time_now_utc)


class SMInvocationEnterMessage(Struct, tag="sm_invocation_enter"):
    uid: str
    iid: str
    timestamp: str = field(default_factory=time_now_utc)


class SMInvocationExitMessage(Struct, tag="sm_invocation_exit"):
    uid: str
    iid: str
    timestamp: str = field(default_factory=time_now_utc)


class SMInvocationErrorMessage(Struct, tag="sm_invocation_error"):
    uid: str
    iid: str
    error_type: str
    error_message: str | None = None
    timestamp: str = field(default_factory=time_now_utc)


class SMInferenceUsageMessage(Struct, tag="sm_inference_usage"):
    uid: str
    iid: str
    usage: dict[str, int]
    timestamp: str = field(default_factory=time_now_utc)


class SMMonadMessage(Struct, tag="sm_monad"):
    uid: str
    iid: str
    body: str
    timestamp: str = field(default_factory=time_now_utc)


class SMInvocationInteractionMessage(Struct, tag="sm_invocation_interaction"):
    uid: str
    iid: str
    event: InteractionEvent
    timestamp: str = field(default_factory=time_now_utc)


class SMInferenceRequestMessage(Struct, tag="sm_inference_request"):
    uid: str
    iid: str
    inference_id: str
    request: str
    timeout: int | None
    timestamp: str = field(default_factory=time_now_utc)


class SMInferenceResponseMessage(Struct, tag="sm_inference_response"):
    uid: str
    iid: str
    inference_id: str
    response: str
    timestamp: str = field(default_factory=time_now_utc)


class SMInferenceErrorMessageMessage(Struct, tag="sm_inference_error"):
    uid: str
    iid: str
    inference_id: str
    error_type: str
    error_message: str
    timestamp: str = field(default_factory=time_now_utc)


SessionManagerMessage = (
    SMInvocationEnterMessage
    | SMInvocationExitMessage
    | SMInvocationErrorMessage
    | SMInferenceErrorMessageMessage
    | SMInferenceRequestMessage
    | SMInferenceResponseMessage
    | SMInvocationInteractionMessage
)


# Messages generated by the agentic function invocation
class MagicExecuteEnterMessage(Struct, tag="mf_execute_enter"):
    uid: str
    iid: str
    code: str
    exec_id: str
    timestamp: str = field(default_factory=time_now_utc)


class MagicExecuteExitMessage(Struct, tag="mf_execute_exit"):
    uid: str
    iid: str
    exec_id: str
    out: str
    stdout: str
    stderr: str
    timestamp: str = field(default_factory=time_now_utc)


class MagicInferenceEnterMessage(Struct, tag="mf_inference_enter"):
    uid: str
    iid: str
    timestamp: str = field(default_factory=time_now_utc)


class MagicInferenceExitMessage(Struct, tag="mf_inference_exit"):
    uid: str
    iid: str
    response: str
    timestamp: str = field(default_factory=time_now_utc)


MagicMessage = (
    MagicExecuteEnterMessage
    | MagicExecuteExitMessage
    | MagicInferenceEnterMessage
    | MagicInferenceExitMessage
)


class MonadStreamMessage(Struct, tag="monad_stream"):
    uid: str
    iid: str
    content: str
    role: Literal['instruction', 'execution', 'code', 'agent']
    timestamp: str = field(default_factory=time_now_utc)


AllServerMessage = (
    SMMonadMessage
    | SMInvocationEnterMessage
    | SMInvocationExitMessage
    | SMInvocationErrorMessage
    | SMInferenceUsageMessage
    | SMInferenceRequestMessage
    | SMInferenceResponseMessage
    | SMInferenceErrorMessageMessage
    | SMInvocationInteractionMessage
    | MagicExecuteEnterMessage
    | MagicExecuteExitMessage
    | MagicInferenceEnterMessage
    | MagicInferenceExitMessage
    | SMCreateAgentMessage
    | SMDestroyAgentMessage
    | MonadStreamMessage
)


_enc = Encoder()


def server_message_to_json(message: AllServerMessage) -> str:
    return _enc.encode(message).decode("utf-8")


def server_message_to_dict(message: AllServerMessage) -> dict[str, Any]:
    # Use to_builtins for direct conversion instead of json.loads(json.dumps(...)) roundtrip
    return to_builtins(message)  # type: ignore[return-value]
