"""
Heracles Cohort Characterization Management

Supports running OHDSI Heracles analyses with full job configuration:
- Custom analysis ID selection
- Job naming and tracking
- Small cell count suppression
- Heracles Heel data quality checks
- Cohort period filtering

Author: Ankur Lohachab
Department of Advanced Computing Sciences, Maastricht University
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, Any, List, Optional
import time


# Standard OHDSI Heracles analysis IDs
# Demographics and summary
DEMO_ANALYSES = [1, 2, 3, 4, 5]

# Condition analyses
CONDITION_ANALYSES = [400, 401, 402, 403, 404, 405, 406, 409, 410, 411, 412, 413]

# Drug analyses
DRUG_ANALYSES = [700, 701, 702, 703, 704, 705, 706, 709, 710, 711, 712, 713]

# Procedure analyses
PROCEDURE_ANALYSES = [600, 601, 602, 603, 604, 605, 606, 609, 610, 611, 612, 613]

# Observation analyses
OBSERVATION_ANALYSES = [800, 801, 802, 803, 804, 805, 806, 809, 810, 811, 812, 813]

# Measurement analyses
MEASUREMENT_ANALYSES = [1800, 1801, 1802, 1803, 1804, 1805, 1806, 1807, 1808, 1809,
                        1810, 1811, 1812, 1813, 1814, 1815, 1816, 1820, 1821, 1830, 1831]

# Visit analyses
VISIT_ANALYSES = [200, 201, 202, 203, 204, 206, 209, 210, 211, 212, 213]

# Default comprehensive analysis set
DEFAULT_ANALYSES = (
    DEMO_ANALYSES +
    CONDITION_ANALYSES +
    DRUG_ANALYSES[:8] +  # Common drug analyses
    PROCEDURE_ANALYSES[:8] +  # Common procedure analyses
    MEASUREMENT_ANALYSES[:10] +  # Common measurement analyses
    VISIT_ANALYSES[:5]  # Common visit analyses
)


@dataclass
class HeraclesJobConfig:
    """Configuration for Heracles characterization job."""

    cohort_definition_ids: List[int]
    source_key: str
    job_name: Optional[str] = None
    analysis_ids: List[int] = field(default_factory=lambda: DEFAULT_ANALYSES.copy())
    small_cell_count: int = 0
    run_heracles_heel: bool = False
    cohort_period_only: bool = False

    def to_dict(self) -> Dict[str, Any]:
        """Convert to WebAPI job submission format."""
        result = {
            "cohortDefinitionIds": self.cohort_definition_ids,
            "sourceKey": self.source_key,
            "analysisIds": self.analysis_ids,
            "smallCellCount": self.small_cell_count,
            "runHeraclesHeel": self.run_heracles_heel,
            "cohortPeriodOnly": self.cohort_period_only
        }

        if self.job_name:
            result["jobName"] = self.job_name

        return result


class HeraclesAnalysisBuilder:
    """Builder for customizing Heracles analysis sets."""

    def __init__(self):
        self.analysis_ids: List[int] = []

    def add_demographics(self) -> HeraclesAnalysisBuilder:
        """Add demographic analyses (person count, age, gender, etc.)."""
        self.analysis_ids.extend(DEMO_ANALYSES)
        return self

    def add_conditions(self) -> HeraclesAnalysisBuilder:
        """Add condition occurrence analyses."""
        self.analysis_ids.extend(CONDITION_ANALYSES)
        return self

    def add_drugs(self, comprehensive: bool = False) -> HeraclesAnalysisBuilder:
        """Add drug exposure analyses."""
        self.analysis_ids.extend(DRUG_ANALYSES if comprehensive else DRUG_ANALYSES[:8])
        return self

    def add_procedures(self, comprehensive: bool = False) -> HeraclesAnalysisBuilder:
        """Add procedure occurrence analyses."""
        self.analysis_ids.extend(PROCEDURE_ANALYSES if comprehensive else PROCEDURE_ANALYSES[:8])
        return self

    def add_measurements(self, comprehensive: bool = False) -> HeraclesAnalysisBuilder:
        """Add measurement analyses."""
        self.analysis_ids.extend(MEASUREMENT_ANALYSES if comprehensive else MEASUREMENT_ANALYSES[:10])
        return self

    def add_visits(self) -> HeraclesAnalysisBuilder:
        """Add visit occurrence analyses."""
        self.analysis_ids.extend(VISIT_ANALYSES)
        return self

    def add_observations(self) -> HeraclesAnalysisBuilder:
        """Add observation analyses."""
        self.analysis_ids.extend(OBSERVATION_ANALYSES)
        return self

    def add_custom(self, *analysis_ids: int) -> HeraclesAnalysisBuilder:
        """Add custom analysis IDs."""
        self.analysis_ids.extend(analysis_ids)
        return self

    def build(self) -> List[int]:
        """Get final analysis ID list (deduplicated)."""
        return sorted(set(self.analysis_ids))


class HeraclesJobManager:
    """
    Manager for Heracles characterization jobs.

    Provides methods to create, submit, and monitor Heracles jobs.
    """

    def __init__(self, client):
        """
        Initialize with OMOPClient instance.

        Args:
            client: OMOPClient instance for WebAPI communication
        """
        self.client = client

    def create_job(
        self,
        cohort_ids: List[int],
        source_key: str,
        job_name: Optional[str] = None,
        analysis_ids: Optional[List[int]] = None,
        small_cell_count: int = 0,
        run_heel: bool = False,
        cohort_period_only: bool = False
    ) -> HeraclesJobConfig:
        """
        Create Heracles job configuration.

        Args:
            cohort_ids: List of cohort definition IDs to characterize
            source_key: Source database key
            job_name: Optional job name for tracking
            analysis_ids: List of analysis IDs to run (defaults to comprehensive set)
            small_cell_count: Suppress cells below this count
            run_heel: Run Heracles Heel data quality checks
            cohort_period_only: Limit analyses to cohort period only

        Returns:
            HeraclesJobConfig instance
        """
        if analysis_ids is None:
            analysis_ids = DEFAULT_ANALYSES.copy()

        return HeraclesJobConfig(
            cohort_definition_ids=cohort_ids,
            source_key=source_key,
            job_name=job_name,
            analysis_ids=analysis_ids,
            small_cell_count=small_cell_count,
            run_heracles_heel=run_heel,
            cohort_period_only=cohort_period_only
        )

    def submit_job(
        self,
        job_config: HeraclesJobConfig,
        poll: bool = False,
        timeout: int = 1800
    ) -> Dict[str, Any]:
        """
        Submit Heracles job to WebAPI.

        Args:
            job_config: Heracles job configuration
            poll: Whether to poll for completion
            timeout: Polling timeout in seconds

        Returns:
            Job response with execution ID and status
        """
        endpoint = "cohortanalysis"
        response = self.client.post(endpoint, job_config.to_dict())

        if not poll:
            return response

        # Poll for completion if requested
        exec_id = response.get('executionId') or response.get('id') or response.get('jobId')

        if not exec_id:
            return response

        return self._poll_job_status(exec_id, timeout)

    def _poll_job_status(
        self,
        execution_id: int,
        timeout: int
    ) -> Dict[str, Any]:
        """Poll job status until completion or timeout."""
        start_time = time.time()
        last_status = None

        while (time.time() - start_time) < timeout:
            try:
                status = self.client.get(f"job/{execution_id}")

                current_status = status.get('status') or status.get('exitStatus')

                if current_status and current_status != last_status:
                    last_status = current_status

                # Terminal states
                if current_status in ('COMPLETED', 'COMPLETE', 'SUCCESS', 'FINISHED'):
                    return status
                elif current_status in ('FAILED', 'ERROR'):
                    raise RuntimeError(f"Heracles job failed: {status}")

                time.sleep(5)

            except Exception:
                # If polling fails, return last known status
                break

        return {'status': 'TIMEOUT', 'executionId': execution_id}

    def get_job_status(self, execution_id: int) -> Dict[str, Any]:
        """Get current status of Heracles job."""
        return self.client.get(f"job/{execution_id}")


def create_standard_job(
    cohort_ids: List[int],
    source_key: str,
    job_name: Optional[str] = None,
    include_demographics: bool = True,
    include_conditions: bool = True,
    include_drugs: bool = True,
    include_procedures: bool = True,
    include_measurements: bool = True
) -> HeraclesJobConfig:
    """
    Create standard Heracles job with common analysis categories.

    Args:
        cohort_ids: Cohort IDs to characterize
        source_key: Source database key
        job_name: Optional job name
        include_demographics: Include demographic analyses
        include_conditions: Include condition analyses
        include_drugs: Include drug analyses
        include_procedures: Include procedure analyses
        include_measurements: Include measurement analyses

    Returns:
        Configured HeraclesJobConfig
    """
    builder = HeraclesAnalysisBuilder()

    if include_demographics:
        builder.add_demographics()
    if include_conditions:
        builder.add_conditions()
    if include_drugs:
        builder.add_drugs()
    if include_procedures:
        builder.add_procedures()
    if include_measurements:
        builder.add_measurements()

    return HeraclesJobConfig(
        cohort_definition_ids=cohort_ids,
        source_key=source_key,
        job_name=job_name,
        analysis_ids=builder.build()
    )
