"""
OHDSI WebAPI Client

HTTP client for interacting with OHDSI WebAPI instances.
Handles authentication, request retries, and error handling.

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

from __future__ import annotations
from typing import Dict, Any, Optional, List
import requests
import time


class OMOPClient:
    """
    Client for OHDSI WebAPI REST endpoints.

    Provides methods for cohort management, concept sets, and data retrieval
    from OMOP Common Data Model databases via WebAPI.
    """

    def __init__(
        self,
        base_url: str,
        timeout: int = 30,
        max_retries: int = 3,
        verify_ssl: bool = True
    ):
        """
        Initialize WebAPI client.

        Args:
            base_url: WebAPI base URL (e.g., "http://host:8080/WebAPI")
            timeout: Request timeout in seconds
            max_retries: Maximum retry attempts for failed requests
            verify_ssl: Whether to verify SSL certificates
        """
        self.base_url = base_url.rstrip('/')
        self.timeout = timeout
        self.max_retries = max_retries
        self.verify_ssl = verify_ssl
        self.session = requests.Session()

    def get(self, endpoint: str, params: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Execute GET request with retry logic.

        Args:
            endpoint: API endpoint path
            params: Query parameters

        Returns:
            JSON response data

        Raises:
            requests.exceptions.RequestException: On request failure
        """
        url = f"{self.base_url}/{endpoint}"

        for attempt in range(self.max_retries):
            try:
                response = self.session.get(
                    url,
                    params=params,
                    timeout=self.timeout,
                    verify=self.verify_ssl
                )
                response.raise_for_status()
                return response.json()

            except requests.exceptions.RequestException as e:
                if attempt == self.max_retries - 1:
                    raise
                # Exponential backoff
                time.sleep(2 ** attempt)

    def post(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute POST request with retry logic.

        Args:
            endpoint: API endpoint path
            data: JSON payload

        Returns:
            JSON response data

        Raises:
            requests.exceptions.RequestException: On request failure
        """
        url = f"{self.base_url}/{endpoint}"

        for attempt in range(self.max_retries):
            try:
                response = self.session.post(
                    url,
                    json=data,
                    timeout=self.timeout,
                    verify=self.verify_ssl,
                    headers={"Content-Type": "application/json"}
                )
                response.raise_for_status()
                return response.json()

            except requests.exceptions.RequestException as e:
                if attempt == self.max_retries - 1:
                    raise
                time.sleep(2 ** attempt)

    def get_sources(self) -> List[Dict[str, Any]]:
        """
        Retrieve available OMOP CDM data sources.

        Returns:
            List of source configurations
        """
        return self.get("source/sources")

    def get_cohort(self, cohort_id: int) -> Dict[str, Any]:
        """
        Retrieve cohort definition by ID.

        Args:
            cohort_id: Cohort definition identifier

        Returns:
            Cohort definition with expression and metadata

        Raises:
            ValueError: If cohort_id is invalid
            RuntimeError: If cohort does not exist or network error
        """
        if not isinstance(cohort_id, int) or cohort_id <= 0:
            raise ValueError(f"Cohort ID must be positive integer, got: {cohort_id}")

        try:
            return self.get(f"cohortdefinition/{cohort_id}")
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise RuntimeError(
                    f"Cohort {cohort_id} does not exist. "
                    f"Use get_sources() to list available cohorts."
                ) from e
            raise RuntimeError(f"Failed to fetch cohort {cohort_id}: {e}") from e
        except requests.exceptions.RequestException as e:
            raise RuntimeError(
                f"Network error fetching cohort {cohort_id}: {e}"
            ) from e

    def create_cohort(self, cohort_definition: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create new cohort definition.

        Args:
            cohort_definition: Cohort definition JSON

        Returns:
            Created cohort with assigned ID

        Raises:
            ValueError: If cohort_definition is invalid
            RuntimeError: If creation fails or network error
        """
        if not isinstance(cohort_definition, dict):
            raise ValueError("Cohort definition must be a dictionary")

        if 'name' not in cohort_definition:
            raise ValueError("Cohort definition must include 'name' field")

        if 'expression' not in cohort_definition:
            raise ValueError("Cohort definition must include 'expression' field")

        try:
            result = self.post("cohortdefinition", cohort_definition)
            if 'id' not in result:
                raise RuntimeError(
                    "Server response missing cohort ID. "
                    f"Response: {result}"
                )
            return result
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 400:
                raise RuntimeError(
                    f"Invalid cohort definition: {e.response.text}"
                ) from e
            elif e.response.status_code == 409:
                raise RuntimeError(
                    f"Cohort with name '{cohort_definition.get('name')}' already exists"
                ) from e
            raise RuntimeError(f"Failed to create cohort: {e}") from e
        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"Network error creating cohort: {e}") from e

    def generate_cohort(
        self,
        cohort_id: int,
        source_key: str
    ) -> Dict[str, Any]:
        """
        Generate cohort on specified data source.

        Args:
            cohort_id: Cohort definition ID
            source_key: Source database key

        Returns:
            Generation job information

        Raises:
            ValueError: If cohort_id or source_key is invalid
            RuntimeError: If generation fails or network error
        """
        if not isinstance(cohort_id, int) or cohort_id <= 0:
            raise ValueError(f"Cohort ID must be positive integer, got: {cohort_id}")

        if not source_key or not isinstance(source_key, str):
            raise ValueError(f"Source key must be non-empty string, got: {source_key}")

        try:
            return self.get(f"cohortdefinition/{cohort_id}/generate/{source_key}")
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise RuntimeError(
                    f"Cohort {cohort_id} or source '{source_key}' does not exist. "
                    f"Use get_sources() to list available sources."
                ) from e
            elif e.response.status_code == 500:
                raise RuntimeError(
                    f"Server error generating cohort {cohort_id} on source '{source_key}'. "
                    f"Check WebAPI logs."
                ) from e
            raise RuntimeError(
                f"Failed to generate cohort {cohort_id} on source '{source_key}': {e}"
            ) from e
        except requests.exceptions.RequestException as e:
            raise RuntimeError(
                f"Network error generating cohort {cohort_id}: {e}"
            ) from e

    def get_generation_status(
        self,
        cohort_id: int,
        source_key: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Check cohort generation status across all sources or specific source.

        Args:
            cohort_id: Cohort definition ID
            source_key: Optional source database key to filter results

        Returns:
            List of generation status objects (one per source)

        Raises:
            ValueError: If cohort_id is invalid
            RuntimeError: If status check fails or network error
        """
        if not isinstance(cohort_id, int) or cohort_id <= 0:
            raise ValueError(f"Cohort ID must be positive integer, got: {cohort_id}")

        try:
            statuses = self.get(f"cohortdefinition/{cohort_id}/info")

            if source_key:
                sources = self.get_sources()
                source_id = next(
                    (s['sourceId'] for s in sources if s['sourceKey'] == source_key),
                    None
                )
                if source_id is None:
                    available = [s['sourceKey'] for s in sources]
                    raise RuntimeError(
                        f"Source '{source_key}' not found. Available: {available}"
                    )
                statuses = [s for s in statuses if s['id']['sourceId'] == source_id]

            return statuses
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise RuntimeError(
                    f"Cohort {cohort_id} does not exist or has not been generated"
                ) from e
            raise RuntimeError(
                f"Failed to get generation status for cohort {cohort_id}: {e}"
            ) from e
        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"Network error getting generation status: {e}") from e

    def get_cohort_results(
        self,
        cohort_id: int,
        source_key: str
    ) -> Dict[str, Any]:
        """
        Retrieve generated cohort results summary.

        Args:
            cohort_id: Cohort definition ID
            source_key: Source database key

        Returns:
            Cohort statistics including personCount and recordCount

        Raises:
            ValueError: If cohort_id or source_key is invalid
            RuntimeError: If results retrieval fails
        """
        if not isinstance(cohort_id, int) or cohort_id <= 0:
            raise ValueError(f"Cohort ID must be positive integer, got: {cohort_id}")

        if not source_key or not isinstance(source_key, str):
            raise ValueError(f"Source key must be non-empty string, got: {source_key}")

        try:
            statuses = self.get_generation_status(cohort_id, source_key)

            if not statuses:
                raise RuntimeError(
                    f"Cohort {cohort_id} has not been generated on source '{source_key}'. "
                    f"Use generate_cohort() first."
                )

            status = statuses[0]

            if status.get('status') == 'FAILED':
                fail_msg = status.get('failMessage', 'Unknown error')
                raise RuntimeError(
                    f"Cohort {cohort_id} generation failed on '{source_key}': {fail_msg}"
                )

            return {
                'status': status.get('status'),
                'personCount': status.get('personCount', 0),
                'recordCount': status.get('recordCount', 0),
                'executionDuration': status.get('executionDuration'),
                'startTime': status.get('startTime'),
                'isValid': status.get('isValid'),
                'isCanceled': status.get('isCanceled'),
                'failMessage': status.get('failMessage')
            }
        except RuntimeError:
            raise
        except Exception as e:
            raise RuntimeError(
                f"Failed to get results for cohort {cohort_id} on '{source_key}': {e}"
            ) from e

    def get_heracles_analyses(
        self,
        cohort_id: int,
        source_key: str
    ) -> List[Dict[str, Any]]:
        """
        Retrieve Heracles characterization analyses list.

        Args:
            cohort_id: Cohort definition ID
            source_key: Source database key

        Returns:
            List of Heracles analyses for the cohort
        """
        endpoint = f"cohortresults/{source_key}/{cohort_id}"
        return self.get(endpoint)

    def get_heracles_report(
        self,
        cohort_id: int,
        source_key: str,
        report_type: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """
        Retrieve specific Heracles report for a cohort.

        Args:
            cohort_id: Cohort definition ID
            source_key: Source database key
            report_type: Type of report (person, condition, drug, procedure,
                        measurement, observation, death, dashboard, etc.)
            refresh: Whether to refresh the report data

        Returns:
            Heracles report data

        Raises:
            ValueError: If parameters are invalid
            RuntimeError: If report retrieval fails
        """
        if not isinstance(cohort_id, int) or cohort_id <= 0:
            raise ValueError(f"Cohort ID must be positive integer, got: {cohort_id}")

        if not source_key or not isinstance(source_key, str):
            raise ValueError(f"Source key must be non-empty string, got: {source_key}")

        if not report_type or not isinstance(report_type, str):
            raise ValueError(f"Report type must be non-empty string, got: {report_type}")

        try:
            endpoint = f"cohortresults/{source_key}/{cohort_id}/{report_type}"
            params = {"refresh": "true"} if refresh else None
            return self.get(endpoint, params=params)
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise RuntimeError(
                    f"Report '{report_type}' not found for cohort {cohort_id} on source '{source_key}'. "
                    f"Make sure Heracles analysis has been run."
                ) from e
            raise RuntimeError(
                f"Failed to get report '{report_type}' for cohort {cohort_id}: {e}"
            ) from e
        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"Network error getting Heracles report: {e}") from e

    def get_heracles_person_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get person-level Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "person", refresh)

    def get_heracles_condition_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get condition occurrence Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "condition", refresh)

    def get_heracles_drug_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get drug exposure Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "drug", refresh)

    def get_heracles_procedure_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get procedure occurrence Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "procedure", refresh)

    def get_heracles_measurement_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get measurement Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "measurement", refresh)

    def get_heracles_observation_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get observation Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "observation", refresh)

    def get_heracles_death_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get death Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "death", refresh)

    def get_heracles_dashboard_report(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get dashboard Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "dashboard", refresh)

    def get_heracles_components_summary(
        self,
        cohort_id: int,
        source_key: str,
        refresh: bool = False
    ) -> Dict[str, Any]:
        """Get components summary Heracles report."""
        return self.get_heracles_report(cohort_id, source_key, "components_summary", refresh)

    def run_heracles(
        self,
        cohort_id: int,
        source_key: str
    ) -> Dict[str, Any]:
        """
        Run Heracles characterization on cohort.

        Args:
            cohort_id: Cohort definition ID
            source_key: Source database key

        Returns:
            Heracles job information
        """
        # Note: Exact endpoint depends on WebAPI version
        # Adjust based on your instance
        endpoint = f"cohortanalysis/{cohort_id}/execute/{source_key}"
        return self.get(endpoint)

    def get_concept_set(self, concept_set_id: int) -> Dict[str, Any]:
        """
        Retrieve concept set definition.

        Args:
            concept_set_id: Concept set identifier

        Returns:
            Concept set with included concepts
        """
        return self.get(f"conceptset/{concept_set_id}")

    def resolve_concept_set(
        self,
        concept_set_expression: List[Dict[str, Any]],
        source_key: str
    ) -> List[int]:
        """
        Resolve concept set expression to concept IDs.

        Args:
            concept_set_expression: Concept set expression items
            source_key: Source database key for vocabulary

        Returns:
            List of resolved concept IDs
        """
        endpoint = f"vocabulary/{source_key}/resolveConceptSetExpression"
        response = self.post(endpoint, concept_set_expression)
        return [item['conceptId'] for item in response.get('items', [])]

    def close(self):
        """Close HTTP session."""
        self.session.close()

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()
