"""
High-Level Data Fetching Functions

Convenience functions for common OMOP data retrieval patterns.

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

from __future__ import annotations
from typing import Dict, Any, List, Optional
from smart_omop.client import OMOPClient


def fetch_cohort_data(
    base_url: str,
    cohort_id: int,
    source_key: str,
    include_results: bool = True,
    **client_kwargs
) -> Dict[str, Any]:
    """
    Fetch complete cohort data including definition and results.

    Args:
        base_url: WebAPI base URL
        cohort_id: Cohort definition ID
        source_key: Source database key
        include_results: Whether to fetch cohort results
        **client_kwargs: Additional OMOPClient arguments

    Returns:
        Dictionary with cohort definition and results

    Example:
        >>> data = fetch_cohort_data(
        ...     "http://host:8080/WebAPI",
        ...     cohort_id=1,
        ...     source_key="MY_CDM"
        ... )
        >>> print(data['definition']['name'])
        >>> print(data['results']['person_count'])
    """
    with OMOPClient(base_url, **client_kwargs) as client:
        # Fetch cohort definition
        definition = client.get_cohort(cohort_id)

        result = {
            "definition": definition,
            "cohort_id": cohort_id,
            "source_key": source_key
        }

        # Fetch results if requested
        if include_results:
            try:
                results = client.get_cohort_results(cohort_id, source_key)
                result["results"] = results
            except Exception as e:
                result["results"] = None
                result["results_error"] = str(e)

        return result


def fetch_concept_sets(
    base_url: str,
    concept_set_ids: List[int],
    **client_kwargs
) -> List[Dict[str, Any]]:
    """
    Fetch multiple concept set definitions.

    Args:
        base_url: WebAPI base URL
        concept_set_ids: List of concept set IDs
        **client_kwargs: Additional OMOPClient arguments

    Returns:
        List of concept set definitions

    Example:
        >>> concept_sets = fetch_concept_sets(
        ...     "http://host:8080/WebAPI",
        ...     [1, 2, 3]
        ... )
        >>> for cs in concept_sets:
        ...     print(cs['name'])
    """
    with OMOPClient(base_url, **client_kwargs) as client:
        concept_sets = []
        for cs_id in concept_set_ids:
            try:
                cs = client.get_concept_set(cs_id)
                concept_sets.append(cs)
            except Exception as e:
                concept_sets.append({
                    "id": cs_id,
                    "error": str(e)
                })

        return concept_sets


def create_and_generate_cohort(
    base_url: str,
    cohort_name: str,
    concept_ids: List[int],
    source_key: str,
    age_min: Optional[int] = None,
    age_max: Optional[int] = None,
    description: str = "",
    **client_kwargs
) -> Dict[str, Any]:
    """
    Create cohort definition and immediately generate on source.

    Args:
        base_url: WebAPI base URL
        cohort_name: Name for cohort
        concept_ids: Condition concept IDs for cohort
        source_key: Source database key for generation
        age_min: Minimum age requirement
        age_max: Maximum age requirement
        description: Cohort description
        **client_kwargs: Additional OMOPClient arguments

    Returns:
        Dictionary with cohort definition and generation status

    Example:
        >>> result = create_and_generate_cohort(
        ...     "http://host:8080/WebAPI",
        ...     "COPD Cohort",
        ...     [255573, 40481087],
        ...     "MY_CDM",
        ...     age_min=40
        ... )
        >>> print(result['cohort_id'])
        >>> print(result['generation_status'])
    """
    from smart_omop.cohort import CohortBuilder

    # Build cohort definition
    builder = CohortBuilder(cohort_name, description)
    builder.with_condition("Primary Condition", concept_ids)

    if age_min is not None or age_max is not None:
        builder.with_age_range(age_min, age_max)

    cohort_def = builder.build()

    # Create and generate
    with OMOPClient(base_url, **client_kwargs) as client:
        # Create cohort
        created = client.create_cohort(cohort_def.to_dict())
        cohort_id = created['id']

        # Generate cohort
        generation = client.generate_cohort(cohort_id, source_key)

        return {
            "cohort_id": cohort_id,
            "definition": created,
            "generation_job": generation,
            "source_key": source_key
        }


def poll_generation_status(
    base_url: str,
    cohort_id: int,
    source_key: str,
    max_wait: int = 300,
    poll_interval: int = 5,
    **client_kwargs
) -> Dict[str, Any]:
    """
    Poll cohort generation status until complete or timeout.

    Args:
        base_url: WebAPI base URL
        cohort_id: Cohort definition ID
        source_key: Source database key
        max_wait: Maximum wait time in seconds
        poll_interval: Polling interval in seconds
        **client_kwargs: Additional OMOPClient arguments

    Returns:
        Final generation status

    Example:
        >>> status = poll_generation_status(
        ...     "http://host:8080/WebAPI",
        ...     cohort_id=1,
        ...     source_key="MY_CDM",
        ...     max_wait=60
        ... )
        >>> print(status['status'])
    """
    import time

    with OMOPClient(base_url, **client_kwargs) as client:
        start_time = time.time()

        while (time.time() - start_time) < max_wait:
            status = client.get_generation_status(cohort_id, source_key)

            # Check if generation complete
            # (exact field depends on WebAPI version)
            if isinstance(status, list):
                for item in status:
                    if item.get('sourceKey') == source_key:
                        if item.get('status') == 'COMPLETE':
                            return item
            elif isinstance(status, dict):
                if status.get('status') == 'COMPLETE':
                    return status

            time.sleep(poll_interval)

        # Timeout
        return {
            "status": "TIMEOUT",
            "elapsed_seconds": time.time() - start_time
        }
