"""Job progress tracking module for OCLAWMA.

Provides real-time progress tracking for long-running jobs with:
- SQLite-backed progress storage
- WebSocket streaming to web UI
- CLI progress monitoring
"""

from __future__ import annotations

import json
import sqlite3
import threading
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

from pydantic import BaseModel, Field, field_validator


class JobProgress(BaseModel):
    """Model representing the progress of a job.

    Attributes:
        job_id: Unique identifier for the job
        progress: Progress percentage (0-100)
        logs: List of log messages with timestamps
        last_updated: ISO format timestamp of last update
        status: Current status of the job (pending, running, completed, failed)
        metadata: Additional job metadata
    """

    job_id: str = Field(..., description="Unique job identifier")
    progress: float = Field(
        default=0.0, ge=0.0, le=100.0, description="Progress percentage (0-100)"
    )
    logs: list[dict[str, Any]] = Field(
        default_factory=list, description="List of log entries with timestamp and message"
    )
    last_updated: str = Field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat(),
        description="ISO format timestamp of last update",
    )
    status: str = Field(
        default="pending", description="Job status: pending, running, completed, failed"
    )
    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional job metadata")

    @field_validator("status")
    @classmethod
    def validate_status(cls, v: str) -> str:
        """Validate status field."""
        allowed = {"pending", "running", "completed", "failed", "cancelled"}
        if v not in allowed:
            raise ValueError(f"Status must be one of: {allowed}")
        return v

    def model_dump_json(self, **kwargs: Any) -> str:
        """Serialize to JSON string."""
        return super().model_dump_json(**kwargs)


@dataclass
class ProgressLogEntry:
    """A single progress log entry."""

    timestamp: str
    message: str
    level: str = "info"


class ProgressStoreError(Exception):
    """Error raised by ProgressStore."""

    pass


class JobNotFoundError(ProgressStoreError):
    """Raised when a job is not found."""

    pass


class ProgressStore:
    """SQLite-backed storage for job progress.

    Thread-safe storage for persisting job progress updates.
    """

    def __init__(self, db_path: Path | str | None = None) -> None:
        """Initialize the progress store.

        Args:
            db_path: Path to SQLite database. If None, uses default location.
        """
        if db_path is None:
            db_path = Path.home() / ".oclawma" / "progress.db"
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        self._local = threading.local()
        self._init_db()

    def _get_connection(self) -> sqlite3.Connection:
        """Get thread-local database connection."""
        if not hasattr(self._local, "connection"):
            self._local.connection = sqlite3.connect(
                str(self.db_path),
                check_same_thread=False,
                isolation_level=None,
            )
            self._local.connection.row_factory = sqlite3.Row
        return self._local.connection

    def _init_db(self) -> None:
        """Initialize database schema."""
        with self._lock:
            conn = self._get_connection()
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS job_progress (
                    job_id TEXT PRIMARY KEY,
                    progress REAL DEFAULT 0,
                    logs TEXT DEFAULT '[]',
                    last_updated TEXT DEFAULT CURRENT_TIMESTAMP,
                    status TEXT DEFAULT 'pending',
                    metadata TEXT DEFAULT '{}'
                )
            """
            )
            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_job_progress_status
                ON job_progress(status)
            """
            )
            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_job_progress_updated
                ON job_progress(last_updated)
            """
            )

    @contextmanager
    def connection(self):
        """Context manager for database connections."""
        conn = None
        try:
            conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
            conn.row_factory = sqlite3.Row
            yield conn
        finally:
            if conn:
                conn.close()

    def save_progress(self, progress: JobProgress) -> None:
        """Save or update job progress.

        Args:
            progress: JobProgress instance to save.
        """
        with self._lock, self.connection() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                INSERT INTO job_progress (job_id, progress, logs, last_updated, status, metadata)
                VALUES (?, ?, ?, ?, ?, ?)
                ON CONFLICT(job_id) DO UPDATE SET
                    progress = excluded.progress,
                    logs = excluded.logs,
                    last_updated = excluded.last_updated,
                    status = excluded.status,
                    metadata = excluded.metadata
                """,
                (
                    progress.job_id,
                    progress.progress,
                    json.dumps(progress.logs),
                    progress.last_updated,
                    progress.status,
                    json.dumps(progress.metadata),
                ),
            )
            conn.commit()

    def get_progress(self, job_id: str) -> JobProgress | None:
        """Get progress for a specific job.

        Args:
            job_id: The job identifier.

        Returns:
            JobProgress instance or None if not found.
        """
        with self._lock, self.connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM job_progress WHERE job_id = ?", (job_id,))
            row = cursor.fetchone()

            if row is None:
                return None

            return JobProgress(
                job_id=row["job_id"],
                progress=row["progress"],
                logs=json.loads(row["logs"]),
                last_updated=row["last_updated"],
                status=row["status"],
                metadata=json.loads(row["metadata"]),
            )

    def list_jobs(
        self, status: str | None = None, limit: int = 100, offset: int = 0
    ) -> list[JobProgress]:
        """List jobs with optional filtering.

        Args:
            status: Filter by status (optional).
            limit: Maximum number of results.
            offset: Number of results to skip.

        Returns:
            List of JobProgress instances.
        """
        with self._lock, self.connection() as conn:
            cursor = conn.cursor()

            if status:
                cursor.execute(
                    """
                        SELECT * FROM job_progress
                        WHERE status = ?
                        ORDER BY last_updated DESC
                        LIMIT ? OFFSET ?
                        """,
                    (status, limit, offset),
                )
            else:
                cursor.execute(
                    """
                        SELECT * FROM job_progress
                        ORDER BY last_updated DESC
                        LIMIT ? OFFSET ?
                        """,
                    (limit, offset),
                )

            rows = cursor.fetchall()
            return [
                JobProgress(
                    job_id=row["job_id"],
                    progress=row["progress"],
                    logs=json.loads(row["logs"]),
                    last_updated=row["last_updated"],
                    status=row["status"],
                    metadata=json.loads(row["metadata"]),
                )
                for row in rows
            ]

    def delete_job(self, job_id: str) -> bool:
        """Delete a job from storage.

        Args:
            job_id: The job identifier.

        Returns:
            True if deleted, False if not found.
        """
        with self._lock, self.connection() as conn:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM job_progress WHERE job_id = ?", (job_id,))
            conn.commit()
            return cursor.rowcount > 0

    def cleanup_old_jobs(self, days: int = 30) -> int:
        """Remove jobs older than specified days.

        Args:
            days: Number of days to keep.

        Returns:
            Number of jobs removed.
        """
        with self._lock, self.connection() as conn:
            cursor = conn.cursor()
            cursor.execute(
                f"""
                    DELETE FROM job_progress
                    WHERE last_updated < datetime('now', '-{days} days')
                    """
            )
            conn.commit()
            return cursor.rowcount


class ProgressTracker:
    """Track progress for a single job.

    Provides methods to update progress, add logs, and retrieve
    current progress state. Automatically persists to storage.
    """

    def __init__(
        self, job_id: str, store: ProgressStore | None = None, auto_save: bool = True
    ) -> None:
        """Initialize progress tracker.

        Args:
            job_id: Unique identifier for this job.
            store: ProgressStore instance. If None, creates default.
            auto_save: Whether to auto-save on each update.
        """
        self.job_id = job_id
        self.store = store or ProgressStore()
        self.auto_save = auto_save
        self._progress = JobProgress(job_id=job_id)
        self._callbacks: list[Callable[[JobProgress], None]] = []
        self._lock = threading.RLock()

        # Save initial state
        if auto_save:
            self.store.save_progress(self._progress)

    def _notify_callbacks(self) -> None:
        """Notify all registered callbacks of progress update."""
        for callback in self._callbacks:
            with suppress(Exception):
                callback(self._progress)

    def on_update(self, callback: Callable[[JobProgress], None]) -> None:
        """Register a callback for progress updates.

        Args:
            callback: Function to call with JobProgress on each update.
        """
        self._callbacks.append(callback)

    def update(self, progress_pct: float) -> None:
        """Update the progress percentage.

        Args:
            progress_pct: Progress percentage (0-100).

        Raises:
            ValueError: If progress is outside 0-100 range.
        """
        if not 0 <= progress_pct <= 100:
            raise ValueError("Progress must be between 0 and 100")

        with self._lock:
            self._progress.progress = progress_pct
            self._progress.last_updated = datetime.now(timezone.utc).isoformat()

            if progress_pct == 0:
                self._progress.status = "pending"
            elif progress_pct == 100:
                self._progress.status = "completed"
            else:
                self._progress.status = "running"

            if self.auto_save:
                self.store.save_progress(self._progress)

            self._notify_callbacks()

    def log(self, message: str, level: str = "info") -> None:
        """Add a log message to the job.

        Args:
            message: Log message text.
            level: Log level (debug, info, warning, error).
        """
        with self._lock:
            log_entry = {
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "message": message,
                "level": level,
            }
            self._progress.logs.append(log_entry)
            self._progress.last_updated = datetime.now(timezone.utc).isoformat()

            if self.auto_save:
                self.store.save_progress(self._progress)

            self._notify_callbacks()

    def set_status(self, status: str) -> None:
        """Set the job status.

        Args:
            status: One of: pending, running, completed, failed, cancelled.

        Raises:
            ValueError: If status is not valid.
        """
        with self._lock:
            self._progress.status = status
            self._progress.last_updated = datetime.now(timezone.utc).isoformat()

            if self.auto_save:
                self.store.save_progress(self._progress)

            self._notify_callbacks()

    def fail(self, error_message: str | None = None) -> None:
        """Mark the job as failed.

        Args:
            error_message: Optional error message to log.
        """
        with self._lock:
            self._progress.status = "failed"
            if error_message:
                self.log(error_message, level="error")
            self._progress.last_updated = datetime.now(timezone.utc).isoformat()

            if self.auto_save:
                self.store.save_progress(self._progress)

            self._notify_callbacks()

    def complete(self, message: str | None = None) -> None:
        """Mark the job as completed.

        Args:
            message: Optional completion message to log.
        """
        with self._lock:
            self._progress.progress = 100.0
            self._progress.status = "completed"
            if message:
                self.log(message, level="info")
            self._progress.last_updated = datetime.now(timezone.utc).isoformat()

            if self.auto_save:
                self.store.save_progress(self._progress)

            self._notify_callbacks()

    def get_progress(self) -> JobProgress:
        """Get the current progress state.

        Returns:
            Current JobProgress instance.
        """
        with self._lock:
            return self._progress.model_copy()

    def set_metadata(self, key: str, value: Any) -> None:
        """Set a metadata value.

        Args:
            key: Metadata key.
            value: Metadata value (must be JSON serializable).
        """
        with self._lock:
            self._progress.metadata[key] = value
            self._progress.last_updated = datetime.now(timezone.utc).isoformat()

            if self.auto_save:
                self.store.save_progress(self._progress)

    def get_metadata(self, key: str, default: Any = None) -> Any:
        """Get a metadata value.

        Args:
            key: Metadata key.
            default: Default value if key not found.

        Returns:
            Metadata value or default.
        """
        return self._progress.metadata.get(key, default)


class ProgressManager:
    """Manage multiple job progress trackers.

    Provides centralized management of job progress tracking
    with support for WebSocket broadcasting.
    """

    def __init__(self, store: ProgressStore | None = None) -> None:
        """Initialize the progress manager.

        Args:
            store: ProgressStore instance. If None, creates default.
        """
        self.store = store or ProgressStore()
        self._trackers: dict[str, ProgressTracker] = {}
        self._lock = threading.RLock()
        self._websocket_callbacks: list[Callable[[str, JobProgress], None]] = []

    def create_tracker(self, job_id: str, auto_save: bool = True) -> ProgressTracker:
        """Create a new progress tracker.

        Args:
            job_id: Unique job identifier.
            auto_save: Whether to auto-save updates.

        Returns:
            New ProgressTracker instance.
        """
        with self._lock:
            tracker = ProgressTracker(job_id=job_id, store=self.store, auto_save=auto_save)

            # Register WebSocket broadcast callback
            if self._websocket_callbacks:

                def broadcast_callback(progress: JobProgress) -> None:
                    for ws_callback in self._websocket_callbacks:
                        with suppress(Exception):
                            ws_callback(job_id, progress)

                tracker.on_update(broadcast_callback)

            self._trackers[job_id] = tracker
            return tracker

    def get_tracker(self, job_id: str) -> ProgressTracker | None:
        """Get an existing tracker.

        Args:
            job_id: Job identifier.

        Returns:
            ProgressTracker or None if not found.
        """
        with self._lock:
            return self._trackers.get(job_id)

    def remove_tracker(self, job_id: str) -> bool:
        """Remove a tracker from memory.

        Args:
            job_id: Job identifier.

        Returns:
            True if removed, False if not found.
        """
        with self._lock:
            if job_id in self._trackers:
                del self._trackers[job_id]
                return True
            return False

    def register_websocket_callback(self, callback: Callable[[str, JobProgress], None]) -> None:
        """Register a WebSocket broadcast callback.

        Args:
            callback: Function(job_id, JobProgress) to call on updates.
        """
        self._websocket_callbacks.append(callback)

    def unregister_websocket_callback(self, callback: Callable[[str, JobProgress], None]) -> None:
        """Unregister a WebSocket callback.

        Args:
            callback: Callback to remove.
        """
        if callback in self._websocket_callbacks:
            self._websocket_callbacks.remove(callback)

    def list_active_jobs(self) -> list[str]:
        """List all active job IDs.

        Returns:
            List of job IDs currently being tracked.
        """
        with self._lock:
            return list(self._trackers.keys())

    def get_job_progress(self, job_id: str) -> JobProgress | None:
        """Get progress for a job from storage.

        Args:
            job_id: Job identifier.

        Returns:
            JobProgress or None if not found.
        """
        # First check active trackers
        tracker = self.get_tracker(job_id)
        if tracker:
            return tracker.get_progress()

        # Fall back to storage
        return self.store.get_progress(job_id)

    def list_jobs(self, status: str | None = None, limit: int = 100) -> list[JobProgress]:
        """List jobs from storage.

        Args:
            status: Filter by status.
            limit: Maximum results.

        Returns:
            List of JobProgress instances.
        """
        return self.store.list_jobs(status=status, limit=limit)
