"""
Job processor for worker.

Single Responsibility: Orchestrate job execution across stages.
"""

import json
from dataclasses import dataclass
from typing import Optional

from ..clients.sqs_client import SQSClient
from .config_downloader import ConfigDownloader
from .error_handler import ErrorHandler
from .stage_registry import StageRegistry
from .status_updater import StatusUpdater
from .visibility_heartbeat import VisibilityHeartbeat


@dataclass
class JobMetadata:
    """Job metadata extracted from SQS message."""

    job_id: str
    job_type: str
    config_s3_path: str
    dataset_name: str
    user_id: str
    config_type: str
    recreate_connections: bool = False


class JobProcessor:
    """
    Process single job from SQS queue.

    Orchestrates job execution flow across stages.
    """

    def __init__(
        self,
        status_updater: StatusUpdater,
        config_downloader: ConfigDownloader,
        stage_registry: StageRegistry,
        error_handler: ErrorHandler,
        sqs_client: SQSClient,
    ):
        """
        Initialize job processor.

        Args:
            status_updater: Status updater for DynamoDB
            config_downloader: Config downloader from S3
            stage_registry: Stage registry for processor lookup
            error_handler: Error handler for failures
            sqs_client: SQS client for visibility heartbeat
        """
        self.status_updater = status_updater
        self.config_downloader = config_downloader
        self.stage_registry = stage_registry
        self.error_handler = error_handler
        self.sqs_client = sqs_client

    def process_job(
        self, sqs_message: dict, receipt_handle: str
    ) -> None:
        """
        Process single job from SQS message.

        Steps:
        1. Extract job metadata from SQS message
        2. Update status to running
        3. Download config from S3
        4. Get stage processor
        5. Execute stage processor (with visibility heartbeat)
        6. Update status to completed

        Args:
            sqs_message: SQS message dict
            receipt_handle: SQS receipt handle for visibility extension

        Raises:
            Exception: If processing fails (prevents message deletion)
        """
        metadata = None
        heartbeat: Optional[VisibilityHeartbeat] = None

        try:
            # Step 1: Extract metadata
            metadata = self._extract_job_metadata(sqs_message)

            print(f"Processing job {metadata.job_id} ({metadata.job_type})")

            # Step 2: Update status to running
            self.status_updater.update_to_running(metadata.job_id)
            self.status_updater.add_processing_stage(
                metadata.job_id, "job_started"
            )

            # Step 3: Download config from S3
            self.status_updater.add_processing_stage(
                metadata.job_id, "downloading_config"
            )
            config_path = self.config_downloader.download_config(
                metadata.config_s3_path, metadata.config_type
            )

            print(f"Config downloaded to: {config_path}")
            self.status_updater.add_processing_stage(
                metadata.job_id, "config_downloaded"
            )

            # Step 4-5: Execute stage with visibility heartbeat
            self.status_updater.add_processing_stage(
                metadata.job_id, f"{metadata.job_type}_stage_starting"
            )

            # Start heartbeat before long-running stage execution
            heartbeat = VisibilityHeartbeat(
                sqs_client=self.sqs_client,
                receipt_handle=receipt_handle,
                heartbeat_interval=240,  # Extend every 4 minutes
                visibility_timeout=300,  # Reset to 5 minutes
            )
            heartbeat.start()

            result = self._execute_stage(
                metadata.job_id,
                metadata.job_type,
                config_path,
                metadata.config_s3_path,
                metadata.recreate_connections,
            )

            print(f"Stage completed: {result}")
            self.status_updater.add_processing_stage(
                metadata.job_id, f"{metadata.job_type}_stage_completed"
            )

            # Step 6: Update status to completed
            self.status_updater.update_to_completed(metadata.job_id, result)

            print(f"Job {metadata.job_id} completed successfully")

        except Exception as e:
            # Handle error and re-raise with stage context
            if metadata:
                print(f"Job {metadata.job_id} failed: {e}")
                # Pass job_type as stage for error enrichment
                self.error_handler.handle_error(
                    metadata.job_id, e, stage=metadata.job_type
                )
            else:
                print(f"Failed to extract job metadata: {e}")
                raise
        finally:
            # Always stop heartbeat on completion or failure
            if heartbeat:
                heartbeat.stop()

    def _extract_job_metadata(self, sqs_message: dict) -> JobMetadata:
        """
        Extract job metadata from SQS message.

        Args:
            sqs_message: SQS message dict

        Returns:
            JobMetadata instance

        Raises:
            ValueError: If message format invalid
        """
        try:
            # Parse message body
            body_str = sqs_message.get("Body")
            if not body_str:
                raise ValueError("Missing message body")

            body = json.loads(body_str)

            return JobMetadata(
                job_id=body["job_id"],
                job_type=body["job_type"],
                config_s3_path=body["config_s3_path"],
                dataset_name=body.get("dataset_name", ""),  # Optional - not used anymore
                user_id=body["user_id"],
                config_type=body.get("config_type", "file"),
                recreate_connections=body.get("recreate_connections", False),
            )
        except (KeyError, json.JSONDecodeError) as e:
            raise ValueError(f"Invalid message format: {e}") from e

    def _execute_stage(self, job_id: str, job_type: str, config_path, config_s3_path: str, recreate_connections: bool = False) -> dict:
        """
        Execute stage processor.

        Args:
            job_id: Job identifier for tracking
            job_type: Type of job (k8s, generate, etc.)
            config_path: Local path to config file/directory
            config_s3_path: S3 path to config (for result metadata)
            recreate_connections: Recreate CloudZero connections (upload stage only)

        Returns:
            Result metadata dict

        Raises:
            RuntimeError: If stage execution fails
        """
        # Get stage processor from registry
        processor = self.stage_registry.get_processor(job_type)

        # Execute processor (delegates to existing stage engines)
        # Pass status_updater for internal tracking
        # Pass recreate_connections only for upload_generated stage
        if job_type == "upload_generated":
            return processor.process(
                config_path,
                config_s3_path=config_s3_path,
                job_id=job_id,
                status_updater=self.status_updater,
                recreate_connections=recreate_connections
            )
        else:
            return processor.process(
                config_path,
                config_s3_path=config_s3_path,
                job_id=job_id,
                status_updater=self.status_updater
            )
