"""
Training Task Repository

Handles training task 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 TrainingTask, TrainingLog, TrainingDocumentLink
from backend.data.repositories.base import BaseRepository

logger = logging.getLogger(__name__)


class TrainingTaskRepository(BaseRepository[TrainingTask]):
    """Repository for training task management.

    Handles:
    - Training task CRUD operations
    - Task status management
    - Training logs
    - Training document links
    """

    def create(
        self,
        admin_token: str,
        name: str,
        task_type: str = "train",
        description: str | None = None,
        config: dict[str, Any] | None = None,
        scheduled_at: datetime | None = None,
        cron_expression: str | None = None,
        is_recurring: bool = False,
        dataset_id: str | None = None,
    ) -> str:
        """Create a new training task.

        Returns:
            Task ID as string
        """
        with get_session_context() as session:
            task = TrainingTask(
                admin_token=admin_token,
                name=name,
                task_type=task_type,
                description=description,
                config=config,
                scheduled_at=scheduled_at,
                cron_expression=cron_expression,
                is_recurring=is_recurring,
                status="scheduled" if scheduled_at else "pending",
                dataset_id=dataset_id,
            )
            session.add(task)
            session.flush()
            return str(task.task_id)

    def get(self, task_id: str) -> TrainingTask | None:
        """Get a training task by ID."""
        with get_session_context() as session:
            result = session.get(TrainingTask, UUID(task_id))
            if result:
                session.expunge(result)
            return result

    def get_by_token(
        self,
        task_id: str,
        admin_token: str | None = None,
    ) -> TrainingTask | None:
        """Get a training task by ID. Token parameter is deprecated."""
        return self.get(task_id)

    def get_paginated(
        self,
        admin_token: str | None = None,
        status: str | None = None,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[TrainingTask], int]:
        """Get paginated training tasks."""
        with get_session_context() as session:
            count_stmt = select(func.count()).select_from(TrainingTask)
            if status:
                count_stmt = count_stmt.where(TrainingTask.status == status)
            total = session.exec(count_stmt).one()

            statement = select(TrainingTask)
            if status:
                statement = statement.where(TrainingTask.status == status)
            statement = statement.order_by(TrainingTask.created_at.desc())
            statement = statement.offset(offset).limit(limit)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results), total

    def get_pending(self) -> list[TrainingTask]:
        """Get pending training tasks ready to run."""
        with get_session_context() as session:
            now = datetime.utcnow()
            statement = select(TrainingTask).where(
                TrainingTask.status.in_(["pending", "scheduled"]),
                (TrainingTask.scheduled_at == None) | (TrainingTask.scheduled_at <= now),
            ).order_by(TrainingTask.created_at)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results)

    def get_running(self) -> TrainingTask | None:
        """Get currently running training task.

        Returns:
            Running task or None if no task is running
        """
        with get_session_context() as session:
            result = session.exec(
                select(TrainingTask)
                .where(TrainingTask.status == "running")
                .order_by(TrainingTask.started_at.desc())
            ).first()
            if result:
                session.expunge(result)
            return result

    def update_status(
        self,
        task_id: str,
        status: str,
        error_message: str | None = None,
        result_metrics: dict[str, Any] | None = None,
        model_path: str | None = None,
    ) -> None:
        """Update training task status."""
        with get_session_context() as session:
            task = session.get(TrainingTask, UUID(task_id))
            if task:
                task.status = status
                task.updated_at = datetime.utcnow()
                if status == "running":
                    task.started_at = datetime.utcnow()
                elif status in ("completed", "failed"):
                    task.completed_at = datetime.utcnow()
                if error_message is not None:
                    task.error_message = error_message
                if result_metrics is not None:
                    task.result_metrics = result_metrics
                if model_path is not None:
                    task.model_path = model_path
                session.add(task)

    def cancel(self, task_id: str) -> bool:
        """Cancel a training task."""
        with get_session_context() as session:
            task = session.get(TrainingTask, UUID(task_id))
            if task and task.status in ("pending", "scheduled"):
                task.status = "cancelled"
                task.updated_at = datetime.utcnow()
                session.add(task)
                return True
            return False

    def add_log(
        self,
        task_id: str,
        level: str,
        message: str,
        details: dict[str, Any] | None = None,
    ) -> None:
        """Add a training log entry."""
        with get_session_context() as session:
            log = TrainingLog(
                task_id=UUID(task_id),
                level=level,
                message=message,
                details=details,
            )
            session.add(log)

    def get_logs(
        self,
        task_id: str,
        limit: int = 100,
        offset: int = 0,
    ) -> list[TrainingLog]:
        """Get training logs for a task."""
        with get_session_context() as session:
            statement = select(TrainingLog).where(
                TrainingLog.task_id == UUID(task_id)
            ).order_by(TrainingLog.created_at.desc()).offset(offset).limit(limit)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results)

    def create_document_link(
        self,
        task_id: UUID,
        document_id: UUID,
        annotation_snapshot: dict[str, Any] | None = None,
    ) -> TrainingDocumentLink:
        """Create a training document link."""
        with get_session_context() as session:
            link = TrainingDocumentLink(
                task_id=task_id,
                document_id=document_id,
                annotation_snapshot=annotation_snapshot,
            )
            session.add(link)
            session.commit()
            session.refresh(link)
            session.expunge(link)
            return link

    def get_document_links(self, task_id: UUID) -> list[TrainingDocumentLink]:
        """Get all document links for a training task."""
        with get_session_context() as session:
            statement = select(TrainingDocumentLink).where(
                TrainingDocumentLink.task_id == task_id
            ).order_by(TrainingDocumentLink.created_at)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results)

    def get_document_training_tasks(self, document_id: UUID) -> list[TrainingDocumentLink]:
        """Get all training tasks that used this document."""
        with get_session_context() as session:
            statement = select(TrainingDocumentLink).where(
                TrainingDocumentLink.document_id == document_id
            ).order_by(TrainingDocumentLink.created_at.desc())

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results)
