"""
Tests for Heracles characterization

Tests Heracles job configuration and management.
"""

import pytest
from smart_omop import (
    HeraclesJobConfig,
    HeraclesAnalysisBuilder,
    create_standard_job,
    DEMO_ANALYSES,
    CONDITION_ANALYSES
)


def test_heracles_job_config():
    """Test Heracles job configuration."""
    config = HeraclesJobConfig(
        cohort_definition_ids=[1, 2],
        source_key="MY_SOURCE",
        job_name="Test Job",
        analysis_ids=[1, 2, 3],
        small_cell_count=5,
        run_heracles_heel=False,
        cohort_period_only=False
    )

    job_dict = config.to_dict()

    assert job_dict['jobName'] == "Test Job"
    assert job_dict['sourceKey'] == "MY_SOURCE"
    assert job_dict['cohortDefinitionIds'] == [1, 2]
    assert job_dict['analysisIds'] == [1, 2, 3]
    assert job_dict['smallCellCount'] == 5
    assert job_dict['runHeraclesHeel'] is False
    assert job_dict['cohortPeriodOnly'] is False


def test_analysis_builder():
    """Test Heracles analysis builder."""
    builder = HeraclesAnalysisBuilder()

    builder.add_demographics()
    builder.add_conditions()
    builder.add_drugs()

    analysis_ids = builder.build()

    assert isinstance(analysis_ids, list)
    assert len(analysis_ids) > 0
    assert all(isinstance(aid, int) for aid in analysis_ids)
    # Should be deduplicated and sorted
    assert analysis_ids == sorted(set(analysis_ids))


def test_analysis_categories():
    """Test predefined analysis categories."""
    assert isinstance(DEMO_ANALYSES, list)
    assert len(DEMO_ANALYSES) > 0

    assert isinstance(CONDITION_ANALYSES, list)
    assert len(CONDITION_ANALYSES) > 0


def test_create_standard_job():
    """Test create_standard_job convenience function."""
    job = create_standard_job(
        cohort_ids=[1],
        source_key="MY_SOURCE",
        job_name="Standard Job",
        include_demographics=True,
        include_conditions=True,
        include_drugs=False
    )

    assert job.job_name == "Standard Job"
    assert job.source_key == "MY_SOURCE"
    assert job.cohort_definition_ids == [1]
    assert len(job.analysis_ids) > 0


def test_custom_analyses():
    """Test adding custom analysis IDs."""
    builder = HeraclesAnalysisBuilder()

    builder.add_custom(9999, 8888, 7777)

    analysis_ids = builder.build()

    assert 9999 in analysis_ids
    assert 8888 in analysis_ids
    assert 7777 in analysis_ids


def test_comprehensive_analyses():
    """Test comprehensive analysis sets."""
    builder = HeraclesAnalysisBuilder()

    builder.add_demographics()
    builder.add_conditions()
    builder.add_drugs(comprehensive=True)
    builder.add_procedures(comprehensive=True)
    builder.add_measurements(comprehensive=True)

    analysis_ids = builder.build()

    # Should have many analyses
    assert len(analysis_ids) > 30
