"""Database operations for training tasks."""

import json
import logging
from datetime import datetime, timezone

import psycopg2
import psycopg2.extras

from shared.config import get_db_connection_string

logger = logging.getLogger(__name__)


class TrainingTaskDB:
    """Read/write training_tasks table."""

    def _connect(self):
        return psycopg2.connect(get_db_connection_string())

    def get_task(self, task_id: str) -> dict | None:
        """Get a single training task by ID."""
        conn = self._connect()
        try:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM training_tasks WHERE task_id = %s",
                    (task_id,),
                )
                return cur.fetchone()
        finally:
            conn.close()

    def get_pending_tasks(self, limit: int = 1) -> list[dict]:
        """Get pending tasks ordered by creation time."""
        conn = self._connect()
        try:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """
                    SELECT * FROM training_tasks
                    WHERE status = 'pending'
                    ORDER BY created_at ASC
                    LIMIT %s
                    """,
                    (limit,),
                )
                return cur.fetchall()
        finally:
            conn.close()

    def update_status(self, task_id: str, status: str) -> None:
        """Update task status with timestamp."""
        conn = self._connect()
        try:
            with conn.cursor() as cur:
                if status == "running":
                    cur.execute(
                        "UPDATE training_tasks SET status = %s, started_at = %s WHERE task_id = %s",
                        (status, datetime.now(timezone.utc), task_id),
                    )
                else:
                    cur.execute(
                        "UPDATE training_tasks SET status = %s WHERE task_id = %s",
                        (status, task_id),
                    )
            conn.commit()
        finally:
            conn.close()

    def complete_task(
        self, task_id: str, model_path: str, metrics: dict
    ) -> None:
        """Mark task as completed with results."""
        conn = self._connect()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    UPDATE training_tasks
                    SET status = 'completed',
                        completed_at = %s,
                        model_path = %s,
                        metrics = %s
                    WHERE task_id = %s
                    """,
                    (
                        datetime.now(timezone.utc),
                        model_path,
                        json.dumps(metrics),
                        task_id,
                    ),
                )
            conn.commit()
        finally:
            conn.close()

    def fail_task(self, task_id: str, error_message: str) -> None:
        """Mark task as failed."""
        conn = self._connect()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    UPDATE training_tasks
                    SET status = 'failed',
                        completed_at = %s,
                        error_message = %s
                    WHERE task_id = %s
                    """,
                    (datetime.now(timezone.utc), error_message[:2000], task_id),
                )
            conn.commit()
        finally:
            conn.close()

    def create_task(self, config: dict) -> str:
        """Create a new training task. Returns task_id."""
        conn = self._connect()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO training_tasks (config)
                    VALUES (%s)
                    RETURNING task_id
                    """,
                    (json.dumps(config),),
                )
                task_id = str(cur.fetchone()[0])
            conn.commit()
            return task_id
        finally:
            conn.close()
