import uuid
from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.models.models import User, Agency


class UserRepository:
    """Repository for user-related database operations."""

    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_by_id(self, user_id: uuid.UUID) -> Optional[User]:
        """Get user by ID."""
        result = await self.session.execute(
            select(User)
            .options(selectinload(User.agency))
            .where(User.id == user_id)
        )
        return result.scalar_one_or_none()

    async def get_by_azure_oid(self, azure_ad_oid: str) -> Optional[User]:
        """Get user by Azure AD Object ID."""
        result = await self.session.execute(
            select(User)
            .options(selectinload(User.agency))
            .where(User.azure_ad_oid == azure_ad_oid)
        )
        return result.scalar_one_or_none()

    async def get_or_create_from_azure(
        self,
        azure_ad_oid: str,
        email: str,
        name: str,
        agency_name: Optional[str] = None
    ) -> User:
        """Get existing user or create new one from Azure AD claims."""
        user = await self.get_by_azure_oid(azure_ad_oid)
        if user:
            return user

        # Create new user
        agency = None
        if agency_name:
            agency = await self.get_or_create_agency(agency_name)

        user = User(
            azure_ad_oid=azure_ad_oid,
            email=email,
            name=name,
            role="basic_user",
            agency_id=agency.id if agency else None,
        )
        self.session.add(user)
        await self.session.flush()
        return user

    async def get_or_create_agency(self, name: str) -> Agency:
        """Get existing agency or create new one."""
        result = await self.session.execute(
            select(Agency).where(Agency.name == name)
        )
        agency = result.scalar_one_or_none()
        if agency:
            return agency

        agency = Agency(name=name)
        self.session.add(agency)
        await self.session.flush()
        return agency

    async def list_all(self) -> list[User]:
        """List all users."""
        result = await self.session.execute(
            select(User).options(selectinload(User.agency))
        )
        return list(result.scalars().all())

    async def update_role(self, user_id: uuid.UUID, role: str) -> Optional[User]:
        """Update user role."""
        user = await self.get_by_id(user_id)
        if user:
            user.role = role
            await self.session.flush()
        return user

    async def list_agencies(self) -> list[Agency]:
        """List all agencies."""
        result = await self.session.execute(select(Agency))
        return list(result.scalars().all())
