"""
Dataset Repository

Handles training dataset operations following Single Responsibility Principle.
"""

import logging
from datetime import datetime
from typing import Any
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 TrainingDataset, DatasetDocument, TrainingTask
from backend.data.repositories.base import BaseRepository

logger = logging.getLogger(__name__)


class DatasetRepository(BaseRepository[TrainingDataset]):
    """Repository for training dataset management.

    Handles:
    - Dataset CRUD operations
    - Dataset status management
    - Dataset document linking
    - Training status tracking
    """

    def create(
        self,
        name: str,
        description: str | None = None,
        train_ratio: float = 0.8,
        val_ratio: float = 0.1,
        seed: int = 42,
    ) -> TrainingDataset:
        """Create a new training dataset."""
        with get_session_context() as session:
            dataset = TrainingDataset(
                name=name,
                description=description,
                train_ratio=train_ratio,
                val_ratio=val_ratio,
                seed=seed,
            )
            session.add(dataset)
            session.commit()
            session.refresh(dataset)
            session.expunge(dataset)
            return dataset

    def get(self, dataset_id: str | UUID) -> TrainingDataset | None:
        """Get a dataset by ID."""
        with get_session_context() as session:
            dataset = session.get(TrainingDataset, UUID(str(dataset_id)))
            if dataset:
                session.expunge(dataset)
            return dataset

    def get_paginated(
        self,
        status: str | None = None,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[TrainingDataset], int]:
        """List datasets with optional status filter."""
        with get_session_context() as session:
            query = select(TrainingDataset)
            count_query = select(func.count()).select_from(TrainingDataset)
            if status:
                query = query.where(TrainingDataset.status == status)
                count_query = count_query.where(TrainingDataset.status == status)
            total = session.exec(count_query).one()
            datasets = session.exec(
                query.order_by(TrainingDataset.created_at.desc()).offset(offset).limit(limit)
            ).all()
            for d in datasets:
                session.expunge(d)
            return list(datasets), total

    def get_active_training_tasks(
        self, dataset_ids: list[str]
    ) -> dict[str, dict[str, str]]:
        """Get active training tasks for datasets.

        Returns a dict mapping dataset_id to {"task_id": ..., "status": ...}
        """
        if not dataset_ids:
            return {}

        valid_uuids = []
        for d in dataset_ids:
            try:
                valid_uuids.append(UUID(d))
            except ValueError:
                logger.warning("Invalid UUID in get_active_training_tasks: %s", d)
                continue

        if not valid_uuids:
            return {}

        with get_session_context() as session:
            statement = select(TrainingTask).where(
                TrainingTask.dataset_id.in_(valid_uuids),
                TrainingTask.status.in_(["pending", "scheduled", "running"]),
            )
            results = session.exec(statement).all()
            return {
                str(t.dataset_id): {"task_id": str(t.task_id), "status": t.status}
                for t in results
            }

    def update_status(
        self,
        dataset_id: str | UUID,
        status: str,
        error_message: str | None = None,
        total_documents: int | None = None,
        total_images: int | None = None,
        total_annotations: int | None = None,
        dataset_path: str | None = None,
    ) -> None:
        """Update dataset status and optional totals."""
        with get_session_context() as session:
            dataset = session.get(TrainingDataset, UUID(str(dataset_id)))
            if not dataset:
                return
            dataset.status = status
            dataset.updated_at = datetime.utcnow()
            if error_message is not None:
                dataset.error_message = error_message
            if total_documents is not None:
                dataset.total_documents = total_documents
            if total_images is not None:
                dataset.total_images = total_images
            if total_annotations is not None:
                dataset.total_annotations = total_annotations
            if dataset_path is not None:
                dataset.dataset_path = dataset_path
            session.add(dataset)
            session.commit()

    def update_training_status(
        self,
        dataset_id: str | UUID,
        training_status: str | None,
        active_training_task_id: str | UUID | None = None,
        update_main_status: bool = False,
    ) -> None:
        """Update dataset training status."""
        with get_session_context() as session:
            dataset = session.get(TrainingDataset, UUID(str(dataset_id)))
            if not dataset:
                return
            dataset.training_status = training_status
            dataset.active_training_task_id = (
                UUID(str(active_training_task_id)) if active_training_task_id else None
            )
            dataset.updated_at = datetime.utcnow()
            if update_main_status and training_status == "completed":
                dataset.status = "trained"
            session.add(dataset)
            session.commit()

    def add_documents(
        self,
        dataset_id: str | UUID,
        documents: list[dict[str, Any]],
    ) -> None:
        """Batch insert documents into a dataset.

        Each dict: {document_id, split, page_count, annotation_count}
        """
        with get_session_context() as session:
            for doc in documents:
                dd = DatasetDocument(
                    dataset_id=UUID(str(dataset_id)),
                    document_id=UUID(str(doc["document_id"])),
                    split=doc["split"],
                    page_count=doc.get("page_count", 0),
                    annotation_count=doc.get("annotation_count", 0),
                )
                session.add(dd)
            session.commit()

    def get_documents(self, dataset_id: str | UUID) -> list[DatasetDocument]:
        """Get all documents in a dataset."""
        with get_session_context() as session:
            results = session.exec(
                select(DatasetDocument)
                .where(DatasetDocument.dataset_id == UUID(str(dataset_id)))
            ).all()
            for r in results:
                session.expunge(r)
            return list(results)

    def delete(self, dataset_id: str | UUID) -> bool:
        """Delete a dataset and its document links."""
        with get_session_context() as session:
            dataset = session.get(TrainingDataset, UUID(str(dataset_id)))
            if not dataset:
                return False
            # Delete associated document links first
            doc_links = session.exec(
                select(DatasetDocument).where(
                    DatasetDocument.dataset_id == UUID(str(dataset_id))
                )
            ).all()
            for link in doc_links:
                session.delete(link)
            session.delete(dataset)
            session.commit()
            return True
