"""ADK Postgres Session Service for GCP Solution Architecture Agent.

Uses google.adk.sessions.SessionService backed by external PostgreSQL database.
"""

import logging
import uuid
from typing import Any, Dict, Optional

from app.adk.compat import Session, SessionService
from app.database import get_db_manager

logger = logging.getLogger(__name__)


class PostgresSessionService(SessionService):
    """ADK Session Service persisting state to PostgreSQL."""

    def __init__(self) -> None:
        self.db_manager = get_db_manager()

    def create_session(self, agent_name: str, session_id: Optional[str] = None) -> Session:
        """Create a new session record in PostgreSQL database."""
        sid = session_id or str(uuid.uuid4())
        session = Session(session_id=sid, agent_name=agent_name, state={})
        self.db_manager.save_session(
            session_id=sid,
            agent_name=agent_name,
            state_data=session.state,
            metadata={"created_via": "PostgresSessionService"},
        )
        logger.info("Created ADK session %s in PostgreSQL database.", sid)
        return session

    def get_session(self, session_id: str) -> Optional[Session]:
        """Fetch session state from PostgreSQL database."""
        sess_data = self.db_manager.get_session(session_id)
        if not sess_data:
            return None

        return Session(
            session_id=sess_data["session_id"],
            agent_name=sess_data["agent_name"],
            state=sess_data.get("state_data", {}),
        )

    def save_session(self, session: Session) -> bool:
        """Update session state in PostgreSQL database."""
        success = self.db_manager.save_session(
            session_id=session.session_id,
            agent_name=session.agent_name,
            state_data=session.state,
        )
        if success:
            logger.info("Updated ADK session %s in PostgreSQL database.", session.session_id)
        return success