"""
Time-shift job monitoring service.

Monitors multiple time-shift jobs until completion.
"""

import time
from typing import List, Set
from ..models.parallel_execution_models import JobSubmission, ProviderResult
from ...api.client import APIClient


class TimeShiftJobMonitor:
    """Monitors multiple time-shift jobs until completion."""

    def __init__(
        self,
        api_client: APIClient,
        poll_interval_sec: int = 5
    ):
        """
        Initialize with API client and poll interval.

        Args:
            api_client: API client for status queries
            poll_interval_sec: Seconds to wait between status checks
        """
        self.api_client = api_client
        self.poll_interval_sec = poll_interval_sec

    def monitor_until_complete(
        self,
        job_submissions: List[JobSubmission]
    ) -> List[ProviderResult]:
        """
        Monitor jobs until all complete or fail.

        Args:
            job_submissions: List of submitted jobs to monitor

        Returns:
            List of ProviderResult in same order as submissions
        """
        job_map = {job.job_id: job for job in job_submissions}
        incomplete = set(job.job_id for job in job_submissions)
        results = {}

        while incomplete:
            completed_ids = self._check_for_completions(incomplete)

            for job_id in completed_ids:
                result = self._fetch_result(job_map[job_id])
                results[job_id] = result
                incomplete.remove(job_id)

            if incomplete:
                time.sleep(self.poll_interval_sec)

        # Return results in original submission order
        return [results[job.job_id] for job in job_submissions]

    def _check_for_completions(
        self,
        job_ids: Set[str]
    ) -> Set[str]:
        """Check which jobs have completed."""
        completed = set()
        for job_id in job_ids:
            status_data = self.api_client.get_job_status(job_id)
            status = status_data.get("status", "unknown")
            if self._is_terminal_status(status):
                completed.add(job_id)
        return completed

    def _is_terminal_status(self, status: str) -> bool:
        """Check if status is terminal (completed or failed)."""
        return status in ("completed", "failed", "success", "error")

    def _fetch_result(
        self,
        submission: JobSubmission
    ) -> ProviderResult:
        """Fetch final result for completed job."""
        job_data = self.api_client.get_job_status(submission.job_id)
        return self._convert_to_provider_result(submission, job_data)

    def _convert_to_provider_result(
        self,
        submission: JobSubmission,
        job_data: dict
    ) -> ProviderResult:
        """Convert API job data to ProviderResult."""
        result = job_data.get("result", {})
        status = job_data.get("status", "unknown")

        return ProviderResult(
            provider=submission.provider,
            job_id=submission.job_id,
            status=status,
            drops_completed=result.get("drops_processed", 0),
            drops_total=result.get("drops_processed", 0) + result.get("drops_failed", 0),
            records_processed=result.get("total_records", 0),
            time_elapsed=result.get("total_time_sec", 0.0),
            error_message=job_data.get("error")
        )
