from sqlalchemy import select

from zbt.domains.agent.memory import CustomerMemory, CustomerMemoryRepository
from zbt.infrastructure.mysql.agent_models import CustomerMemoryRecord
from zbt.infrastructure.mysql.repositories import SessionFactory


class SqlAlchemyCustomerMemoryRepository(CustomerMemoryRepository):
    def __init__(self, session_factory: SessionFactory) -> None:
        self._session_factory = session_factory

    def save(self, memory: CustomerMemory) -> None:
        with self._session_factory() as session:
            session.merge(
                CustomerMemoryRecord(
                    id=memory.id,
                    owner_id=memory.owner_id,
                    category=memory.category,
                    content=memory.content,
                    status=memory.status,
                    created_at=memory.created_at,
                    updated_at=memory.updated_at,
                )
            )
            session.commit()

    def get(self, memory_id: str) -> CustomerMemory | None:
        with self._session_factory() as session:
            record = session.get(CustomerMemoryRecord, memory_id)
            return self._to_domain(record) if record is not None else None

    def list_active(self, owner_id: str) -> list[CustomerMemory]:
        with self._session_factory() as session:
            records = session.scalars(
                select(CustomerMemoryRecord)
                .where(
                    CustomerMemoryRecord.owner_id == owner_id,
                    CustomerMemoryRecord.status == "ACTIVE",
                )
                .order_by(CustomerMemoryRecord.updated_at.desc())
            ).all()
            return [self._to_domain(record) for record in records]

    @staticmethod
    def _to_domain(record: CustomerMemoryRecord) -> CustomerMemory:
        return CustomerMemory(
            id=record.id,
            owner_id=record.owner_id,
            category=record.category,
            content=record.content,
            status=record.status,
            created_at=record.created_at,
            updated_at=record.updated_at,
        )
