from sqlmodel import select
from typing import List, Optional
from kilostar.core.postgres_database.model.workflow import EventRecord
from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession


class EventDatabase:
    def __init__(self, async_session_maker: async_sessionmaker[AsyncSession]):
        self.async_session_maker = async_session_maker

    async def upsert_event(self, trace_id: str, event_data_json: str) -> EventRecord:
        async with self.async_session_maker() as session:
            statement = select(EventRecord).where(EventRecord.trace_id == trace_id)
            results = await session.execute(statement)
            record = results.scalar_one_or_none()
            if record:
                record.event_data_json = event_data_json
            else:
                record = EventRecord(trace_id=trace_id, event_data_json=event_data_json)
                session.add(record)
            await session.commit()
            await session.refresh(record)
            return record

    async def get_event(self, trace_id: str) -> Optional[EventRecord]:
        async with self.async_session_maker() as session:
            statement = select(EventRecord).where(EventRecord.trace_id == trace_id)
            results = await session.execute(statement)
            return results.scalar_one_or_none()

    async def get_all_events(self) -> List[EventRecord]:
        async with self.async_session_maker() as session:
            statement = select(EventRecord)
            results = await session.execute(statement)
            return results.scalars().all()

    async def delete_event(self, trace_id: str) -> bool:
        async with self.async_session_maker() as session:
            statement = select(EventRecord).where(EventRecord.trace_id == trace_id)
            results = await session.execute(statement)
            record = results.scalar_one_or_none()
            if record:
                await session.delete(record)
                await session.commit()
                return True
            return False