from __future__ import annotations

import base64
import json
import uuid
from datetime import datetime

from pydantic import BaseModel, Field, field_serializer, field_validator, model_validator


class Envelope[T: BaseModel | None](BaseModel):
    id: str | None = None
    """The id of the message"""
    ts: datetime | None = Field(default_factory=datetime.now)
    """The timestamp of the message"""
    correlation_id: str | None = None
    """The correlation id of the message"""
    reply_to: str | None = None
    """If the message is a reply, this is the id of the original message"""
    data: T | None = Field(default=None, description="The data of the message")
    """The data of the message"""
    type: str | None = Field(default=None, description="The type of the message")
    """The type of the message"""

    @model_validator(mode="after")
    def _fill_defaults(self) -> "Envelope[T]":
        if self.type is None:
            msg_type = getattr(self.data, "__msg_type__", None)
            if msg_type is None:
                # Fallback (fine for internal systems, less great for long-term compatibility)
                msg_type = f"{self.data.__class__.__module__}.{self.data.__class__.__name__}"
            self.type = msg_type
        return self

    @staticmethod
    def from_data[TData: BaseModel](data: TData, id: str | bool | None = None) -> Envelope[TData]:
        if id is True:
            id = str(uuid.uuid4())

        elif id is False:
            id = None

        return Envelope(data=data, id=id)


# Test envelope class
class TestMessage(BaseModel):
    __msg_type__ = "test"
    test_field: str = Field(..., description="The test field")


# Test float array message with a custom serializer
class FloatArray(BaseModel):
    __msg_type__ = "float_array"
    array: list[float] = Field(..., description="The array of floats")

    @field_serializer("array")
    def serialize_array(self, array: list[float]) -> str:
        # Base64 encode the array
        return base64.b64encode(json.dumps(array).encode()).decode()

    @field_validator("array", mode="before")
    @classmethod
    def deserialize_array(cls, v):
        # already decoded?
        if isinstance(v, list):
            return v

        # encoded string / bytes?
        if isinstance(v, (str, bytes, bytearray)):
            try:
                raw = v.encode("ascii") if isinstance(v, str) else bytes(v)
                decoded = base64.b64decode(raw)
                return json.loads(decoded.decode("utf-8"))
            except Exception as e:
                raise ValueError("Invalid base64-encoded float array") from e


# from faststream import FastStream
# from faststream.nats import NatsBroker

# broker = NatsBroker("nats://localhost:4222")

# app = FastStream(broker)


# @broker.subscriber("test")
# async def test_handler(env: Envelope[FloatArray]):
#     print("got:", env.data)


# @broker.subscriber("test2")
# async def test_handler2(id: str | None, data: FloatArray):
#     print("got2:", data)


# @app.after_startup
# async def publish_once():
#     o = Envelope.from_data(FloatArray(array=[1.0, 2.0, 3.0]))
#     await broker.publish(o.model_dump(mode="json"), subject="test")  # <-- important
#     await broker.publish(o.model_dump(mode="json"), subject="test2")  # <-- important


# async def main():
#     await app.run()  # blocking


# if __name__ == "__main__":
#     o = Envelope.from_data(data=FloatArray(array=[1.0, 2.0, 3.0]))
#     print(o.model_dump_json())

#     # Test the deserialization
#     env = Envelope[FloatArray].model_validate_json(o.model_dump_json())
#     print(env)

#     asyncio.run(main())
