"""
ByteAstra backend — MongoDB document wrapper models.
Behaves like SQLAlchemy ORM classes for seamless API compatibility.
"""
import enum
from datetime import datetime, timezone


class MessageRole(str, enum.Enum):
    user = "user"
    assistant = "assistant"
    system = "system"


class Message:
    """Wraps a MongoDB message document."""
    def __init__(self, doc: dict):
        self.id = doc["_id"]
        self.session_id = doc["session_id"]
        self.role = MessageRole(doc["role"])
        self.content = doc["content"]
        self.citations_json = doc.get("citations_json")
        self.is_grounded = doc.get("is_grounded")
        self.created_at = doc.get("created_at") or datetime.now(timezone.utc)


class Session:
    """Wraps a MongoDB session document and its messages."""
    def __init__(self, doc: dict, messages: list[Message] = None):
        self.id = doc["_id"]
        self.domain = doc["domain"]
        self.title = doc["title"]
        self.user_id = doc.get("user_id")
        self.subscription_tier = doc.get("subscription_tier", "free")
        self.created_at = doc.get("created_at") or datetime.now(timezone.utc)
        self.updated_at = doc.get("updated_at") or datetime.now(timezone.utc)
        self.messages = messages or []