"""
Fine-Tune Pool Repository

Manages the fine-tune pool: accumulated user-reported failure documents
for incremental model improvement.
"""

import logging
from datetime import datetime
from uuid import UUID

from sqlalchemy import func
from sqlmodel import select

from backend.data.database import get_session_context
from backend.data.admin_models import FineTunePoolEntry
from backend.data.repositories.base import BaseRepository

logger = logging.getLogger(__name__)


class FineTunePoolRepository(BaseRepository[FineTunePoolEntry]):
    """Repository for fine-tune pool management."""

    def add_document(
        self,
        document_id: str | UUID,
        added_by: str | None = None,
        reason: str | None = None,
    ) -> FineTunePoolEntry:
        """Add a document to the fine-tune pool."""
        with get_session_context() as session:
            entry = FineTunePoolEntry(
                document_id=UUID(str(document_id)),
                added_by=added_by,
                reason=reason,
            )
            session.add(entry)
            session.commit()
            session.refresh(entry)
            session.expunge(entry)
            return entry

    def get_entry(self, entry_id: str | UUID) -> FineTunePoolEntry | None:
        """Get a pool entry by ID."""
        with get_session_context() as session:
            entry = session.get(FineTunePoolEntry, UUID(str(entry_id)))
            if entry:
                session.expunge(entry)
            return entry

    def get_by_document(self, document_id: str | UUID) -> FineTunePoolEntry | None:
        """Get pool entry for a document."""
        with get_session_context() as session:
            result = session.exec(
                select(FineTunePoolEntry).where(
                    FineTunePoolEntry.document_id == UUID(str(document_id))
                )
            ).first()
            if result:
                session.expunge(result)
            return result

    def get_paginated(
        self,
        verified_only: bool = False,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[FineTunePoolEntry], int]:
        """List pool entries with pagination."""
        with get_session_context() as session:
            query = select(FineTunePoolEntry)
            count_query = select(func.count()).select_from(FineTunePoolEntry)
            if verified_only:
                query = query.where(FineTunePoolEntry.is_verified == True)
                count_query = count_query.where(FineTunePoolEntry.is_verified == True)
            total = session.exec(count_query).one()
            entries = session.exec(
                query.order_by(FineTunePoolEntry.created_at.desc())
                .offset(offset)
                .limit(limit)
            ).all()
            for e in entries:
                session.expunge(e)
            return list(entries), total

    def get_pool_count(self, verified_only: bool = True) -> int:
        """Get count of entries in the pool."""
        with get_session_context() as session:
            query = select(func.count()).select_from(FineTunePoolEntry)
            if verified_only:
                query = query.where(FineTunePoolEntry.is_verified == True)
            return session.exec(query).one()

    def get_all_document_ids(self, verified_only: bool = True) -> list[UUID]:
        """Get all document IDs in the pool."""
        with get_session_context() as session:
            query = select(FineTunePoolEntry.document_id)
            if verified_only:
                query = query.where(FineTunePoolEntry.is_verified == True)
            results = session.exec(query).all()
            return list(results)

    def verify_entry(
        self,
        entry_id: str | UUID,
        verified_by: str | None = None,
    ) -> FineTunePoolEntry | None:
        """Mark a pool entry as verified."""
        with get_session_context() as session:
            entry = session.get(FineTunePoolEntry, UUID(str(entry_id)))
            if not entry:
                return None
            entry.is_verified = True
            entry.verified_at = datetime.utcnow()
            entry.verified_by = verified_by
            session.add(entry)
            session.commit()
            session.refresh(entry)
            session.expunge(entry)
            return entry

    def remove_entry(self, entry_id: str | UUID) -> bool:
        """Remove an entry from the pool."""
        with get_session_context() as session:
            entry = session.get(FineTunePoolEntry, UUID(str(entry_id)))
            if not entry:
                return False
            session.delete(entry)
            session.commit()
            return True
