from __future__ import annotations

import asyncio
from pyroute import Route
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional, Hashable, Dict, Self
from nemoria.cryptography import uniqID


"""
Core protocol types for Nemoria.

This module defines three core types exchanged or managed by the transport:

- `Action`: the set of verb tokens sent over the wire (e.g., "GET", "SET")
- `Frame`: a single protocol message (request/response) with optional
           correlation via `reply_to`
- `Connection`: lightweight metadata for a connected peer (id/host/port),
                plus its bound asyncio streams (reader/writer)
"""

__all__ = ("JSON", "Action", "Frame", "Connection")

# JSON-safe type alias (keys must be JSON-serializable, typically strings)
JSON = Dict[Hashable, Any]


class Action(Enum):
    """
    Enumeration of all supported frame actions.

    Each member represents a wire-level verb in the client-server protocol.
    These are used to express key-value CRUD and control operations.
    The string token carried on the wire is the enum's value (e.g., "GET").
    """

    GET = "GET"
    ALL = "ALL"
    SET = "SET"
    DELETE = "DELETE"
    DROP = "DROP"
    PURGE = "PURGE"
    SAVE = "SAVE"
    PING = "PING"
    HANDSHAKE = "HANDSHAKE"
    SHUTDOWN = "SHUTDOWN"

    @classmethod
    def from_str(cls, name: str) -> Action:
        """
        Convert a token into an `Action`.

        Notes:
            The lookup is performed against the enum **name** (e.g., "GET"),
            which matches the value strings used on the wire in this protocol.

        Args:
            name: Action name token (e.g., "GET", "PING").

        Returns:
            The corresponding `Action` enum.

        Raises:
            ValueError: If the token is not a valid action name.
        """
        try:
            return cls[name]
        except KeyError:
            raise ValueError("invalid frame action")


@dataclass(frozen=True)
class Frame:
    """
    Protocol frame exchanged between client and server.

    A frame encapsulates a single request or response. Correlation between a
    response and its originating request is achieved via the `reply_to` field,
    which may itself contain a serialized `Frame`. In practice, a minimal stub
    (e.g., only `id` and possibly `action`) is sufficient for correlation, but
    `serialize()` preserves the full structure by default.

    Fields:
        id:       Unique identifier (auto-generated by default).
        action:   The protocol verb for this frame (string token on the wire).
        route:    Optional path for addressing a key/namespace.
        value:    Optional payload (must be JSON-serializable by your transport).
        reply_to: Optional original request frame for correlation.
    """

    id: str = field(default_factory=uniqID)
    action: Optional[Action] = None
    route: Optional[Route] = None
    value: Optional[Any] = None
    reply_to: Optional[Frame] = None

    def __eq__(self, other: object) -> bool:
        """Compare frames by their unique IDs."""
        if not isinstance(other, Frame):
            return NotImplemented
        return self.id == other.id

    def __hash__(self) -> int:
        """Hash a frame by its ID (usable in sets and as dict keys)."""
        return hash(self.id)

    def serialize(self) -> JSON:
        """
        Convert the frame into a JSON-safe dict for transmission.

        Notes:
            - `action` is serialized as its string token (e.g., "GET").
            - `route` is serialized via `Route.serialize()`.
            - `reply_to` is serialized **recursively**.

        Returns:
            A JSON-serializable dictionary representing the frame.
        """
        return {
            "id": self.id,
            "action": None if self.action is None else self.action.value,
            "route": None if self.route is None else self.route.serialize(),
            "value": self.value,
            "reply_to": None if self.reply_to is None else self.reply_to.serialize(),
        }

    @classmethod
    def deserialize(cls, obj: JSON) -> Self:
        """
        Reconstruct a `Frame` from its serialized dict form.

        Args:
            obj: JSON-safe dict produced by `serialize()`.

        Returns:
            A `Frame` instance reconstructed from the payload.

        Raises:
            TypeError: If `obj` is not a mapping-like structure.
            KeyError: If required keys are missing from `obj`.
            ValueError: If `action` or `route` fail to decode.
        """
        if not isinstance(obj, dict):
            raise TypeError("Frame JSON must be a dict.")
        return cls(
            id=obj["id"],
            action=None if obj["action"] is None else Action.from_str(obj["action"]),
            route=None if obj["route"] is None else Route.deserialize(obj["route"]),
            value=obj.get("value"),
            reply_to=(
                None
                if obj.get("reply_to") is None
                else Frame.deserialize(obj["reply_to"])
            ),
        )


@dataclass(frozen=True)
class Connection:
    """
    Represents an active peer connection.

    Contains basic peer metadata (`id`, `host`, `port`) and the bound asyncio
    stream endpoints used by the transport layer.

    Warning:
        `reader` and `writer` are live stream objects and are **not**
        serializable; `serialize()` returns only metadata.
    """

    id: str
    host: str
    port: int
    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter

    def __repr__(self) -> str:
        """Compact, log-friendly representation."""
        return f"Connection(Port: {self.port})"

    def serialize(self) -> JSON:
        """
        Convert the connection metadata into a JSON-safe dict.

        Returns:
            Dict with `id`, `host`, and `port`. Stream endpoints are excluded.
        """
        return {"id": self.id, "host": self.host, "port": self.port}

    @classmethod
    def deserialize(
        cls, obj: JSON, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> Self:
        """
        Reconstruct a `Connection` from serialized metadata plus streams.

        Args:
            obj: Dict with keys "id", "host", "port".
            reader: Bound asyncio `StreamReader`.
            writer: Bound asyncio `StreamWriter`.

        Returns:
            A `Connection` instance.

        Raises:
            TypeError: If `obj` is not a dict.
            KeyError: If required keys are missing from `obj`.
        """
        if not isinstance(obj, dict):
            raise TypeError("Connection JSON must be a dict.")
        return cls(obj["id"], obj["host"], obj["port"], reader, writer)
