"""Scenario repository with specific methods."""

from typing import Optional, List
from uuid import UUID
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update
from enum import Enum

from src.models.scenario import Scenario
from src.repositories.base import BaseRepository


class ScenarioStatus(str, Enum):
    """Scenario status enum."""

    DRAFT = "draft"
    RUNNING = "running"
    COMPLETED = "completed"
    ARCHIVED = "archived"


class ScenarioRepository(BaseRepository[Scenario]):
    """Repository for Scenario model with specific methods."""

    def __init__(self):
        super().__init__(Scenario)

    async def get_by_name(self, db: AsyncSession, name: str) -> Optional[Scenario]:
        """Get scenario by name."""
        result = await db.execute(select(Scenario).where(Scenario.name == name))
        return result.scalar_one_or_none()

    async def list_by_status(
        self, db: AsyncSession, status: ScenarioStatus, skip: int = 0, limit: int = 100
    ) -> List[Scenario]:
        """List scenarios by status."""
        return await self.get_multi(db, status=status.value, skip=skip, limit=limit)

    async def list_by_region(
        self, db: AsyncSession, region: str, skip: int = 0, limit: int = 100
    ) -> List[Scenario]:
        """List scenarios by region."""
        return await self.get_multi(db, region=region, skip=skip, limit=limit)

    async def update_status(
        self, db: AsyncSession, scenario_id: UUID, new_status: ScenarioStatus
    ) -> Optional[Scenario]:
        """Update scenario status."""
        result = await db.execute(
            update(Scenario)
            .where(Scenario.id == scenario_id)
            .values(status=new_status.value)
            .returning(Scenario)
        )
        await db.commit()
        return result.scalar_one_or_none()

    async def increment_total_requests(
        self, db: AsyncSession, scenario_id: UUID, increment: int = 1
    ) -> None:
        """Atomically increment total_requests counter."""
        await db.execute(
            update(Scenario)
            .where(Scenario.id == scenario_id)
            .values(total_requests=Scenario.total_requests + increment)
        )
        await db.commit()

    async def update_total_cost(
        self, db: AsyncSession, scenario_id: UUID, new_cost: float
    ) -> None:
        """Update total cost estimate."""
        await db.execute(
            update(Scenario)
            .where(Scenario.id == scenario_id)
            .values(total_cost_estimate=new_cost)
        )
        await db.commit()


# Singleton instance
scenario_repository = ScenarioRepository()
