from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Any, Dict, Mapping, Optional

from agno.utils.log import logger


@dataclass
class AgentSession:
    """Agent Session that is stored in the database"""

    # Session UUID
    session_id: str
    # ID of the agent that this session is associated with
    agent_id: Optional[str] = None
    # ID of the user interacting with this agent
    user_id: Optional[str] = None
    # Agent Memory
    memory: Optional[Dict[str, Any]] = None
    # Agent Data: agent_id, name and model
    agent_data: Optional[Dict[str, Any]] = None
    # Session Data: session_name, session_state, images, videos, audio
    session_data: Optional[Dict[str, Any]] = None
    # Extra Data stored with this agent
    extra_data: Optional[Dict[str, Any]] = None
    # The unix timestamp when this session was created
    created_at: Optional[int] = None
    # The unix timestamp when this session was last updated
    updated_at: Optional[int] = None

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def monitoring_data(self) -> Dict[str, Any]:
        # Google Gemini adds a "parts" field to the messages, which is not serializable
        # If the provider is Google, remove the "parts" from the messages
        if self.agent_data is not None:
            if self.agent_data.get("model", {}).get("provider") == "Google" and self.memory is not None:
                # Remove parts from runs' response messages
                if "runs" in self.memory:
                    for _run in self.memory["runs"]:
                        if "response" in _run and "messages" in _run["response"]:
                            for m in _run["response"]["messages"]:
                                if isinstance(m, dict):
                                    m.pop("parts", None)

                # Remove parts from top-level memory messages
                if "messages" in self.memory:
                    for m in self.memory["messages"]:
                        if isinstance(m, dict):
                            m.pop("parts", None)

        monitoring_data = asdict(self)
        return monitoring_data

    def telemetry_data(self) -> Dict[str, Any]:
        return {
            "model": self.agent_data.get("model") if self.agent_data else None,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> Optional[AgentSession]:
        if data is None or data.get("session_id") is None:
            logger.warning("AgentSession is missing session_id")
            return None
        return cls(
            session_id=data.get("session_id"),  # type: ignore
            agent_id=data.get("agent_id"),
            user_id=data.get("user_id"),
            memory=data.get("memory"),
            agent_data=data.get("agent_data"),
            session_data=data.get("session_data"),
            extra_data=data.get("extra_data"),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at"),
        )
