"""
MedSynth OMOP Data Source Integration

Provides data source implementation for MedSynth-generated OMOP CDM data (CSV format).
Allows using synthetic OMOP data generated by the medsynth package for testing and
development without requiring a full OHDSI WebAPI deployment.

Requires medsynth package: pip install medsynth

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

from __future__ import annotations
from pathlib import Path
from typing import Dict, Any, List, Optional
import json


class MedSynthOMOPSource:
    """
    OMOP data source for MedSynth-generated CSV files.

    Provides cohort-like access to synthetic OMOP data without requiring WebAPI.
    """

    def __init__(self, data_directory: str):
        """
        Initialize MedSynth data source.

        Args:
            data_directory: Path to directory containing MedSynth OMOP CSV files

        Raises:
            ValueError: If data_directory is invalid or does not exist
        """
        if not data_directory:
            raise ValueError("Data directory path is required")

        self.data_dir = Path(data_directory)

        if not self.data_dir.exists():
            raise ValueError(
                f"Data directory not found: {data_directory}. "
                f"Use generate_synthetic_omop() to create OMOP data."
            )

        if not self.data_dir.is_dir():
            raise ValueError(f"Path is not a directory: {data_directory}")

        # Expected OMOP tables
        self.table_files = {
            'person': 'person.csv',
            'condition_occurrence': 'condition_occurrence.csv',
            'drug_exposure': 'drug_exposure.csv',
            'procedure_occurrence': 'procedure_occurrence.csv',
            'measurement': 'measurement.csv',
            'observation': 'observation.csv',
            'visit_occurrence': 'visit_occurrence.csv',
            'death': 'death.csv'
        }

    def get_table_path(self, table_name: str) -> Optional[Path]:
        """Get path to OMOP table CSV file."""
        filename = self.table_files.get(table_name.lower())
        if not filename:
            return None

        path = self.data_dir / filename
        return path if path.exists() else None

    def load_table(self, table_name: str) -> List[Dict[str, Any]]:
        """
        Load OMOP table as list of dicts.

        Args:
            table_name: Name of OMOP table (person, condition_occurrence, etc.)

        Returns:
            List of row dictionaries

        Raises:
            ValueError: If table_name is invalid
            RuntimeError: If table file cannot be read
        """
        import csv

        if not table_name:
            raise ValueError("Table name is required")

        if table_name not in self.table_files:
            available = list(self.table_files.keys())
            raise ValueError(
                f"Unknown table '{table_name}'. Available tables: {available}"
            )

        path = self.get_table_path(table_name)
        if not path:
            raise RuntimeError(
                f"Table file for '{table_name}' not found in {self.data_dir}. "
                f"Expected file: {self.table_files[table_name]}"
            )

        try:
            rows = []
            with open(path, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    rows.append(dict(row))
            return rows
        except Exception as e:
            raise RuntimeError(
                f"Failed to read table '{table_name}' from {path}: {e}"
            ) from e

    def get_person_count(self) -> int:
        """Get total number of persons in dataset."""
        persons = self.load_table('person')
        return len(persons)

    def get_condition_counts(self) -> Dict[str, int]:
        """Get counts of conditions by concept ID."""
        conditions = self.load_table('condition_occurrence')
        counts: Dict[str, int] = {}

        for cond in conditions:
            concept_id = cond.get('condition_concept_id', '0')
            counts[concept_id] = counts.get(concept_id, 0) + 1

        return counts

    def filter_by_condition(
        self,
        concept_ids: List[int],
        include_descendants: bool = False
    ) -> List[str]:
        """
        Filter persons by condition concept IDs.

        Args:
            concept_ids: List of condition concept IDs
            include_descendants: Whether to include descendant concepts (not implemented for CSV)

        Returns:
            List of person_ids matching criteria
        """
        conditions = self.load_table('condition_occurrence')
        concept_id_strs = {str(cid) for cid in concept_ids}

        person_ids = set()
        for cond in conditions:
            if cond.get('condition_concept_id') in concept_id_strs:
                person_ids.add(cond.get('person_id'))

        return sorted(person_ids)

    def filter_by_age_gender(
        self,
        person_ids: List[str],
        min_age: Optional[int] = None,
        max_age: Optional[int] = None,
        gender_concept_ids: Optional[List[int]] = None
    ) -> List[str]:
        """
        Filter person IDs by age and gender criteria.

        Args:
            person_ids: List of person IDs to filter
            min_age: Minimum age (optional)
            max_age: Maximum age (optional)
            gender_concept_ids: List of gender concept IDs (optional)

        Returns:
            Filtered list of person IDs
        """
        persons = self.load_table('person')
        person_id_set = set(person_ids)

        filtered = []
        for person in persons:
            pid = person.get('person_id')

            if pid not in person_id_set:
                continue

            # Check age
            year_of_birth = person.get('year_of_birth')
            if year_of_birth and (min_age or max_age):
                from datetime import datetime
                current_year = datetime.now().year
                age = current_year - int(year_of_birth)

                if min_age and age < min_age:
                    continue
                if max_age and age > max_age:
                    continue

            # Check gender
            if gender_concept_ids:
                gender_id = person.get('gender_concept_id')
                if gender_id and int(gender_id) not in gender_concept_ids:
                    continue

            filtered.append(pid)

        return filtered

    def create_cohort_summary(
        self,
        concept_ids: List[int],
        min_age: Optional[int] = None,
        max_age: Optional[int] = None,
        gender_concept_ids: Optional[List[int]] = None
    ) -> Dict[str, Any]:
        """
        Create cohort summary matching WebAPI response format.

        Args:
            concept_ids: Condition concept IDs
            min_age: Minimum age
            max_age: Maximum age
            gender_concept_ids: Gender concept IDs

        Returns:
            Cohort summary dict compatible with fetch_cohort_data
        """
        # Filter by condition
        person_ids = self.filter_by_condition(concept_ids)

        # Apply demographics filter
        if min_age or max_age or gender_concept_ids:
            person_ids = self.filter_by_age_gender(
                person_ids, min_age, max_age, gender_concept_ids
            )

        # Get demographics breakdown
        persons = self.load_table('person')
        person_map = {p.get('person_id'): p for p in persons}

        gender_counts: Dict[int, int] = {}
        age_sum = 0
        age_count = 0

        for pid in person_ids:
            person = person_map.get(pid)
            if not person:
                continue

            # Gender
            gid = person.get('gender_concept_id')
            if gid:
                gender_counts[int(gid)] = gender_counts.get(int(gid), 0) + 1

            # Age
            year_of_birth = person.get('year_of_birth')
            if year_of_birth:
                from datetime import datetime
                age = datetime.now().year - int(year_of_birth)
                age_sum += age
                age_count += 1

        avg_age = age_sum / age_count if age_count > 0 else 0

        return {
            'source_type': 'medsynth',
            'personCount': len(person_ids),
            'recordCount': len(person_ids),
            'status': 'COMPLETE',
            'demographics': {
                'gender_distribution': gender_counts,
                'average_age': round(avg_age, 1),
                'age_count': age_count
            },
            'person_ids': person_ids[:100],  # First 100 for privacy
            'concept_ids': concept_ids
        }

    def export_to_bundle(self) -> Dict[str, Any]:
        """Export data source metadata as bundle for model cards."""
        return {
            'data_source': {
                'type': 'MedSynth Synthetic OMOP',
                'directory': str(self.data_dir),
                'person_count': self.get_person_count(),
                'available_tables': [
                    table for table in self.table_files.keys()
                    if self.get_table_path(table)
                ]
            }
        }


def load_from_medsynth_directory(
    data_directory: str,
    concept_ids: List[int],
    min_age: Optional[int] = None,
    max_age: Optional[int] = None,
    gender_concept_ids: Optional[List[int]] = None
) -> Dict[str, Any]:
    """
    Convenience function to load cohort data from MedSynth OMOP directory.

    Args:
        data_directory: Path to MedSynth OMOP CSV directory
        concept_ids: Condition concept IDs to filter
        min_age: Minimum age filter
        max_age: Maximum age filter
        gender_concept_ids: Gender concept IDs to filter

    Returns:
        Cohort summary dict compatible with fetch_cohort_data
    """
    source = MedSynthOMOPSource(data_directory)

    return source.create_cohort_summary(
        concept_ids=concept_ids,
        min_age=min_age,
        max_age=max_age,
        gender_concept_ids=gender_concept_ids
    )


def generate_synthetic_omop(
    output_directory: str,
    num_persons: int = 1000,
    condition_concept_ids: Optional[List[int]] = None
) -> str:
    """
    Generate synthetic OMOP data using MedSynth package.

    Requires: pip install medsynth

    Args:
        output_directory: Directory to save CSV files
        num_persons: Number of synthetic persons to generate
        condition_concept_ids: Condition concepts to include

    Returns:
        Path to output directory

    Raises:
        ImportError: If medsynth package not installed
    """
    try:
        from medsynth import generate_omop_cdm
    except ImportError:
        raise ImportError(
            "MedSynth package not installed. Install with: pip install medsynth"
        )

    # Generate synthetic OMOP CDM data
    output_path = Path(output_directory)
    output_path.mkdir(parents=True, exist_ok=True)

    generate_omop_cdm(
        output_dir=str(output_path),
        num_patients=num_persons,
        condition_concepts=condition_concept_ids or [255573]  # Default to COPD
    )

    return str(output_path)
