python/src/server/utils/progress/progress_tracker.py

"""
Progress Tracker Utility

Tracks operation progress in memory for HTTP polling access.
"""

from datetime import datetime
from typing import Any, cast

from ...config.logfire_config import safe_logfire_error, safe_logfire_info


class ProgressTracker:
    """
    Utility class for tracking progress updates in memory.
    State can be accessed via HTTP polling endpoints.
    """

    # Class-level storage for all progress states
    _progress_states: dict[str, dict[str, Any]] = {}

    def __init__(self, progress_id: str, operation_type: str = "crawl"):
        """
        Initialize the progress tracker.

        Args:
            progress_id: Unique progress identifier
            operation_type: Type of operation (crawl, upload, etc.)
        """
        self.progress_id = progress_id
        self.operation_type = operation_type
        self.state: dict[str, Any] = {
            "progress_id": progress_id,
            "type": operation_type,  # Store operation type for progress model selection
            "start_time": datetime.now().isoformat(),
            "status": "initializing",
            "progress": 0,
            "logs": [],
        }
        # Store in class-level dictionary
        ProgressTracker._progress_states[progress_id] = self.state

    @classmethod
    def get_progress(cls, progress_id: str) -> dict[str, Any] | None:
        """Get progress state by ID."""
        return cls._progress_states.get(progress_id)

    @classmethod
    def clear_progress(cls, progress_id: str) -> None:
        """Remove progress state from memory."""
        if progress_id in cls._progress_states:
            del cls._progress_states[progress_id]

    async def start(self, initial_data: dict[str, Any] | None = None):
        """
        Start progress tracking with initial data.

        Args:
            initial_data: Optional initial data to include
        """
        self.state["status"] = "starting"
        self.state["start_time"] = datetime.now().isoformat()

        if initial_data:
            self.state.update(initial_data)

        self._update_state()
        safe_logfire_info(f"Progress tracking started | progress_id={self.progress_id} | type={self.operation_type}")

    async def update(self, status: str, progress: int, log: str, **kwargs):
        """
        Update progress with status, progress, and log message.

        Args:
            status: Current status (analyzing, crawling, processing, etc.)
            progress: Progress value (0-100)
            log: Log message describing current operation
            **kwargs: Additional data to include in update
        """
        # CRITICAL: Never allow progress to go backwards
        current_progress = cast(int, self.state.get("progress", 0))
        new_progress = min(100, max(0, progress))  # Ensure 0-100

        # Only update if new progress is greater than or equal to current
        # (equal allows status updates without progress regression)
        if new_progress < current_progress:
            safe_logfire_info(
                f"Progress backwards prevented: {current_progress}% -> {new_progress}% | "
                f"progress_id={self.progress_id} | status={status}"
            )
            # Keep the higher progress value
            actual_progress = current_progress
        else:
            actual_progress = new_progress

        self.state.update(
            {
                "status": status,
                "progress": actual_progress,
                "log": log,
                "timestamp": datetime.now().isoformat(),
            }
        )

        # Add log entry
        if "logs" not in self.state:
            self.state["logs"] = []

        logs_list = cast(list[dict[str, Any]], self.state["logs"])
        logs_list.append(
            {
                "timestamp": datetime.now().isoformat(),
                "message": log,
                "status": status,
                "progress": progress,
            }
        )

        # Add any additional data
        for key, value in kwargs.items():
            self.state[key] = value

        self._update_state()

    async def complete(self, completion_data: dict[str, Any] | None = None):
        """
        Mark progress as completed with optional completion data.

        Args:
            completion_data: Optional data about the completed operation
        """
        self.state["status"] = "completed"
        self.state["progress"] = 100
        self.state["end_time"] = datetime.now().isoformat()

        if completion_data:
            self.state.update(completion_data)

        # Calculate duration
        if "start_time" in self.state:
            start_str = cast(str, self.state["start_time"])
            end_str = cast(str, self.state["end_time"])
            start = datetime.fromisoformat(start_str)
            end = datetime.fromisoformat(end_str)
            duration = (end - start).total_seconds()
            self.state["duration"] = str(duration)  # Convert to string for Pydantic model
            self.state["duration_formatted"] = self._format_duration(duration)

        self._update_state()
        safe_logfire_info(
            f"Progress completed | progress_id={self.progress_id} | type={self.operation_type} | duration={self.state.get('duration_formatted', 'unknown')}"
        )

    async def error(self, error_message: str, error_details: dict[str, Any] | None = None):
        """
        Mark progress as failed with error information.

        Args:
            error_message: Error message
            error_details: Optional additional error details
        """
        self.state.update(
            {
                "status": "error",
                "error": error_message,
                "error_time": datetime.now().isoformat(),
            }
        )

        if error_details:
            self.state["error_details"] = error_details

        self._update_state()
        safe_logfire_error(
            f"Progress error | progress_id={self.progress_id} | type={self.operation_type} | error={error_message}"
        )

    async def update_batch_progress(self, current_batch: int, total_batches: int, batch_size: int, message: str):
        """
        Update progress for batch operations.

        Args:
            current_batch: Current batch number (1-based)
            total_batches: Total number of batches
            batch_size: Size of each batch
            message: Progress message
        """
        progress_val = int((current_batch / total_batches) * 100)
        await self.update(
            status="processing_batch",
            progress=progress_val,
            log=message,
            current_batch=current_batch,
            total_batches=total_batches,
            batch_size=batch_size,
        )

    async def update_crawl_stats(self, processed_pages: int, total_pages: int, current_url: str | None = None):
        """
        Update crawling statistics.

        Args:
            processed_pages: Number of pages processed
            total_pages: Total pages to process
            current_url: Currently processing URL
        """
        progress_val = int((processed_pages / max(total_pages, 1)) * 100)
        log = f"Processing page {processed_pages}/{total_pages}"
        if current_url:
            log += f": {current_url}"

        await self.update(
            status="crawling",
            progress=progress_val,
            log=log,
            processed_pages=processed_pages,
            total_pages=total_pages,
            current_url=current_url,
        )

    async def update_storage_progress(self, chunks_stored: int, total_chunks: int, operation: str = "storing"):
        """
        Update document storage progress.

        Args:
            chunks_stored: Number of chunks stored
            total_chunks: Total chunks to store
            operation: Storage operation description
        """
        progress_val = int((chunks_stored / max(total_chunks, 1)) * 100)
        await self.update(
            status="document_storage",
            progress=progress_val,
            log=f"{operation}: {chunks_stored}/{total_chunks} chunks",
            chunks_stored=chunks_stored,
            total_chunks=total_chunks,
        )

    async def update_code_extraction_progress(
        self,
        documents_processed: int,
        total_documents: int,
        code_blocks_found: int,
        total_summaries: int,
        current_file: str,
    ):
        """
        Update code extraction progress.

        Args:
            documents_processed: Number of documents processed
            total_documents: Total documents to process
            code_blocks_found: Number of code blocks found
            total_summaries: Total summaries to process
            current_file: Currently processing file
        """
        progress_val = int((documents_processed / max(total_documents, 1)) * 100)
        await self.update(
            status="code_extraction",
            progress=progress_val,
            log=f"Processing document {documents_processed}/{total_documents}",
            documents_processed=documents_processed,
            total_documents=total_documents,
            code_blocks_found=code_blocks_found,
            total_summaries=total_summaries,
            current_file=current_file,
        )

    def _update_state(self):
        """Update progress state in memory storage."""
        # Update the class-level dictionary
        ProgressTracker._progress_states[self.progress_id] = self.state

        safe_logfire_info(
            f"📊 [PROGRESS] Updated {self.operation_type} | ID: {self.progress_id} | "
            f"Status: {self.state.get('status')} | Progress: {self.state.get('progress')}%"
        )

    def _format_duration(self, seconds: float) -> str:
        """Format duration in seconds to human-readable string."""
        if seconds < 60:
            return f"{seconds:.1f} seconds"
        elif seconds < 3600:
            minutes = seconds / 60
            return f"{minutes:.1f} minutes"
        else:
            hours = seconds / 3600
            return f"{hours:.1f} hours"

    def get_state(self) -> dict[str, Any]:
        """Get current progress state."""
        return self.state.copy()