from ibis.expr.types.relations import Table

from typing import Union
import ibis
import copy
import logging

logger = logging.getLogger(__name__)

from phenex.util.serialization.to_dict import to_dict


class PhenexTable:
    """
    Phenex provides certain table types on which it knows how to operate. For instance, Phenex implements a CodeTable, which is an event table containing codes. Phenex has abstracted operations for each table type. For instance, given a CodeTable, Phenex knows how to filter this table based on the presence of codes within that table. Phenex doesn't care if the code table is actually a diagnosis code table or a procedure code table or a medication code table.

    In onboarding a new data model to Phenex, the tables must be mapped into Phenex table types by subclassing the appropriate PhenexTable. When subclassing a PhenexTable, you must define:

        1. COLUMN_MAPPING: a mapping of the input table columns to the fields on the chosen PhenexTable type (e.g. 'CD' maps to 'CODE' in a CodeTable).
        2. JOIN_KEYS: if you want to use the autojoin functionality of PhenexTable, you must specify what keys to use to join pairs of tables
        3. PATHS: if you want to use the autojoin functionality of PhenexTable for more complex joins, you must specify join paths to take to get from one table to another

    Note that for each table type, there are REQUIRED_FIELDS, i.e., fields that MUST be defined for Phenex to work with such a table and KNOWN_FIELDS, i.e., fields that Phenex internally understands what to do with (there is a Phenotype that knows how to work with that field). For instance, in a PhenexPersonTable, one MUST define PERSON_ID, but DATE_OF_BIRTH is an optional field that PhenEx can process if given and transform into AGE. These are fixed for each table type and should not be overridden.

    DEFAULT_MAPPING can map to either a single column (string) or multiple columns (list):
    - String value: Creates a column as a direct copy/rename
      Example: {"PERSON_ID": "PERSON_ID"} copies PERSON_ID column
    - List value: Creates a column as a coalesce of multiple columns (first non-null value)
      Example: {"EVENT_DATE": ["STARTDATETIME", "RECORDEDDATETIME"]} creates EVENT_DATE using
      STARTDATETIME if available, falling back to RECORDEDDATETIME if STARTDATETIME is null

    JOIN_KEYS and PATHS Documentation:

    JOIN_KEYS defines direct relationships between tables. The key is the CLASS NAME of the target table,
    and the value is a list of join keys. Each join key can be:
    - A string: symmetric join (column has same name in both tables)
    - A 2-element tuple/list: asymmetric join (left_col, right_col) with different names

    PATHS defines multi-hop join paths. The key is the CLASS NAME of the final target table,
    and the value is a list of CLASS NAMES for intermediate tables to traverse.

    IMPORTANT: JOIN_KEYS should be defined symmetrically - if TableA can join to TableB,
    then TableB should also define how to join back to TableA.

    Example 1: Symmetric joins (same column names)
    ```python
    class DummyConditionOccurrenceTable(CodeTable):
        NAME_TABLE = "DIAGNOSIS"
        JOIN_KEYS = {
            "DummyPersonTable": ["PERSON_ID"],  # Join using PERSON_ID in both tables
            "DummyEncounterTable": ["PERSON_ID", "ENCID"],  # Compound join: both keys must match
        }
        PATHS = {
            "DummyVisitDetailTable": ["DummyEncounterTable"]  # To reach VisitDetail, go through Encounter
        }

    class DummyEncounterTable(PhenexTable):
        NAME_TABLE = "ENCOUNTER"
        JOIN_KEYS = {
            "DummyPersonTable": ["PERSON_ID"],
            "DummyConditionOccurrenceTable": ["PERSON_ID", "ENCID"],  # Symmetric!
            "DummyVisitDetailTable": ["PERSON_ID", "VISITID"],
        }

    class DummyVisitDetailTable(PhenexTable):
        NAME_TABLE = "VISIT"
        JOIN_KEYS = {
            "DummyPersonTable": ["PERSON_ID"],
            "DummyEncounterTable": ["PERSON_ID", "VISITID"],  # Symmetric!
        }
    ```

    Example 2: Asymmetric joins (different column names)
    ```python
    class EventTable(CodeTable):
        NAME_TABLE = "EVENT"
        JOIN_KEYS = {
            "EventMappingTable": [("ID", "EVENTID")],  # EventTable.ID joins to EventMappingTable.EVENTID
        }
        PATHS = {
            "ConceptTable": ["EventMappingTable"],
        }
        DEFAULT_MAPPING = {
            "PERSON_ID": "PERSON_ID",
            "ID": "ID",  # Must include ID in mapping for it to exist
        }

    class EventMappingTable(PhenexTable):
        NAME_TABLE = "EVENT_MAPPING"
        JOIN_KEYS = {
            "EventTable": [("EVENTID", "ID")],  # Symmetric: reverse the tuple
            "ConceptTable": [("CONCEPTID", "ID")],  # Maps to ConceptTable.ID
        }
        DEFAULT_MAPPING = {
            "EVENTID": "EVENTID",
            "CONCEPTID": "CONCEPTID",
        }

    class ConceptTable(CodeTable):
        NAME_TABLE = "CONCEPT"
        JOIN_KEYS = {
            "EventMappingTable": [("ID", "CONCEPTID")],  # Symmetric: reverse the tuple
        }
        DEFAULT_MAPPING = {
            "ID": "ID",
            "CODE": "CONCEPT_CODE",
            "CODE_TYPE": "VOCABULARY_ID",
        }
    ```

    Example 3: Mixed symmetric and asymmetric joins
    ```python
    class PatientEventTable(CodeTable):
        JOIN_KEYS = {
            "EventMappingTable": [
                "PERSON_ID",  # Symmetric: PERSON_ID in both tables
                ("EVENT_ID", "EVENTID")  # Asymmetric: different column names
            ],
        }
    ```

    In all examples:
    - Symmetric relationships use strings: ["COLUMN_NAME"]
    - Asymmetric relationships use tuples: [("LEFT_COL", "RIGHT_COL")]
    - Compound joins use multiple elements: ["COL1", "COL2"] or [("L1", "R1"), ("L2", "R2")]
    - All relationships should be symmetric (both tables define the join)
    - ALL join columns must be in DEFAULT_MAPPING for them to exist in the mapped table
    """

    NAME_TABLE = "PHENEX_TABLE"  # name of table in the database
    JOIN_KEYS = {}  # dict: class name -> List[phenex column names]
    KNOWN_FIELDS = []  # List[phenex column names]
    DEFAULT_MAPPING = {}  # dict: input column name -> phenex column name
    PATHS = {}  # dict: table class name -> List[other table class names]
    REQUIRED_FIELDS = list(DEFAULT_MAPPING.keys())

    def __init__(self, table, name=None, column_mapping={}):
        """
        Instantiate a PhenexTable, possibly overriding NAME_TABLE and COLUMN_MAPPING.
        """

        if not isinstance(table, Table):
            raise TypeError(
                f"Cannot instantiatiate {self.__class__.__name__} from {type(table)}. Must be ibis Table."
            )

        self.NAME_TABLE = name or self.NAME_TABLE

        self.column_mapping = self._get_column_mapping(column_mapping)
        self._table = table.mutate(
            **self._resolve_column_mapping(table, self.column_mapping)
        )

        for key in self.REQUIRED_FIELDS:
            try:
                getattr(self._table, key)
            except AttributeError:
                raise ValueError(f"Required field {key} not defined in COLUMN_MAPPING.")

        self._add_phenotype_table_relationship()

    def _add_phenotype_table_relationship(self):
        self.JOIN_KEYS["PhenotypeTable"] = ["PERSON_ID"]

    def _get_column_mapping(self, column_mapping=None):
        column_mapping = column_mapping or {}
        # Only validate fields explicitly passed in column_mapping parameter
        # DEFAULT_MAPPING is defined by the class itself and should be trusted
        # This allows join keys and other auxiliary fields to be in DEFAULT_MAPPING
        # without requiring them to be in KNOWN_FIELDS
        for key in column_mapping.keys():
            if key not in self.KNOWN_FIELDS:
                raise ValueError(
                    f"Unknown mapped field {key} --> {column_mapping[key]} for f{type(self)}."
                )
        default_mapping = copy.deepcopy(self.DEFAULT_MAPPING)
        default_mapping.update(column_mapping)
        return default_mapping

    def _resolve_column_mapping(self, table, column_mapping):
        """
        Convert raw column mapping (strings/lists) to ibis expressions for use in mutate().

        String values become direct column references: table[col].
        List values become coalesce expressions over the listed columns.
        Date columns in a coalesce list are cast to timestamp for consistent typing.
        """
        processed_mapping = {}
        for key, value in column_mapping.items():
            if isinstance(value, list):
                # Coalesce multiple columns - first non-null value wins
                # Cast date columns to timestamp for consistent typing
                cols = []
                for col in value:
                    col_ref = table[col]
                    col_type = str(col_ref.type())
                    if col_type.startswith("date") and not col_type.startswith(
                        "timestamp"
                    ):
                        col_ref = col_ref.cast("timestamp")
                    cols.append(col_ref)
                processed_mapping[key] = ibis.coalesce(*cols)
            else:
                # Single column mapping - create column reference
                processed_mapping[key] = table[value]
        return processed_mapping

    def __getattr__(self, name):
        # pass all attributes on to underlying table
        return getattr(self._table, name)

    def __getitem__(self, key):
        return self._table[key]

    @property
    def table(self):
        return self._table

    def join(self, other: "PhenexTable", *args, domains=None, **kwargs):
        """
        The join method performs a join of PhenexTables, using autojoin functionality if Phenex is able to find the table types specified in PATHS.
        """
        if isinstance(other, Table):
            return type(self)(self.table.join(other, *args, **kwargs))

        if not isinstance(other, PhenexTable):
            raise TypeError(f"Expected a PhenexTable instance, got {type(other)}")
        if len(args):
            # if user specifies join keys and join type, simply perform join as specified
            return type(self)(self.table.join(other.table, *args, **kwargs))

        # Do an autojoin by finding a path from the left to the right table and sequentially joining as necessary
        # joined table is the sequentially joined table
        # current table is the table for the left join in the current iteration
        joined_table = current_left_table = self
        logger.debug(
            f"Starting autojoin from {self.__class__.__name__} to {other.__class__.__name__}"
        )

        for right_table_class_name in self._find_path(other):
            # get the next right table
            right_table_search_results = [
                v
                for k, v in domains.items()
                if v.__class__.__name__ == right_table_class_name
            ]
            logger.debug(
                f"Searching for {right_table_class_name} in domains: {list(domains.keys())}"
            )
            logger.debug(
                f"Found {len(right_table_search_results)} matches for {right_table_class_name}"
            )

            if len(right_table_search_results) != 1:
                raise ValueError(
                    f"Unable to find unqiue {right_table_class_name} required to join {other.__class__.__name__}"
                )
            right_table = right_table_search_results[0]
            print(
                f"\tJoining : {current_left_table.__class__.__name__} to {right_table.__class__.__name__}"
            )

            # join keys are defined by the left table; in theory should enforce symmetry
            join_keys = current_left_table.JOIN_KEYS[right_table_class_name]

            # Build join predicate(s) - supports symmetric and asymmetric joins
            # Symmetric: ["COLUMN"] or ["COL1", "COL2"] - same column names in both tables
            # Asymmetric: [("LEFT_COL", "RIGHT_COL")] - different column names
            # Mixed: ["COL1", ("LEFT_COL", "RIGHT_COL")]
            predicates = []
            for join_key in join_keys:
                if isinstance(join_key, str):
                    # Symmetric: column exists in both tables with same name
                    predicates.append(joined_table[join_key] == right_table[join_key])
                elif isinstance(join_key, (tuple, list)) and len(join_key) == 2:
                    # Asymmetric: (left_col, right_col) - different column names
                    left_col, right_col = join_key
                    predicates.append(joined_table[left_col] == right_table[right_col])
                else:
                    raise ValueError(
                        f"Invalid join key format: {join_key}. Must be either a string or a 2-element tuple/list."
                    )

            # Combine all predicates with AND
            if len(predicates) == 1:
                join_predicate = predicates[0]
            else:
                join_predicate = predicates[0]
                for pred in predicates[1:]:
                    join_predicate = join_predicate & pred

            columns = list(set(joined_table.columns + right_table.columns))
            # subset columns, making sure to set type of table to the very left table (self)
            joined_table = type(self)(
                joined_table.join(right_table, join_predicate, **kwargs).select(columns)
            )
            current_left_table = right_table
        return joined_table

    def mutate(self, *args, **kwargs):
        return type(self)(self.table.mutate(*args, **kwargs), name=self.NAME_TABLE)

    def _find_path(self, other):
        start_name = self.__class__.__name__
        end_name = other.__class__.__name__

        logger.debug(f"Finding path from {start_name} to {end_name}")

        # first see if direct connection
        try:
            join_keys = self.JOIN_KEYS[end_name]
            logger.debug(
                f"Found direct connection: {start_name} -> {end_name} using keys {join_keys}"
            )
            return [end_name]
        except KeyError:
            logger.debug(
                f"No direct connection found in JOIN_KEYS for {start_name} -> {end_name}"
            )
            try:
                path = self.PATHS[end_name]
                full_path = path + [end_name]
                logger.debug(
                    f"Found path in PATHS: {start_name} -> {' -> '.join(full_path)}"
                )
                return full_path
            except KeyError:
                logger.error(f"No path found for {start_name} -> {end_name}")
                logger.debug(
                    f"Available JOIN_KEYS for {start_name}: {list(self.JOIN_KEYS.keys())}"
                )
                logger.debug(
                    f"Available PATHS for {start_name}: {list(self.PATHS.keys())}"
                )
                raise ValueError(
                    f"Cannot autojoin {start_name} --> {end_name}. Please specify join path in PATHS."
                )

    def filter(self, expr):
        """
        Filter the table by an Ibis Expression or using a PhenExFilter.
        """
        input_columns = self.columns
        if isinstance(expr, ibis.expr.types.Expr) or isinstance(expr, list):
            filtered_table = self.table.filter(expr)
        else:
            filtered_table = expr.filter(self)

        return type(self)(
            filtered_table.select(input_columns),
            name=self.NAME_TABLE,
            column_mapping=self.column_mapping,
        )

    @classmethod
    def to_dict(cls) -> dict:
        """
        Serialize the PhenexTable class configuration (not the data).

        This serializes the class-level attributes that define the table mapping,
        but not the actual ibis table data which cannot be serialized.

        Returns:
            dict: Class configuration including NAME_TABLE, JOIN_KEYS, DEFAULT_MAPPING, etc.
        """
        return {
            "__table_class__": cls.__name__,
            "__module__": cls.__module__,
            "NAME_TABLE": cls.NAME_TABLE,
            "JOIN_KEYS": cls.JOIN_KEYS,
            "KNOWN_FIELDS": cls.KNOWN_FIELDS,
            "DEFAULT_MAPPING": cls.DEFAULT_MAPPING,
            "PATHS": cls.PATHS,
            "REQUIRED_FIELDS": cls.REQUIRED_FIELDS,
        }

    @classmethod
    def from_dict(cls, data: dict):
        """
        Reconstruct a PhenexTable class reference from serialized data.

        Note: This returns the class itself, not an instance, since we cannot
        reconstruct the actual table data without a database connection.

        Args:
            data: Serialized class configuration

        Returns:
            The PhenexTable subclass
        """
        # The class should already exist in the module, just return it
        return cls


class PhenexPersonTable(PhenexTable):
    NAME_TABLE = "PERSON"
    JOIN_KEYS = {
        "CodeTable": ["PERSON_ID"],
        "PhenexVisitOccurrenceTable": ["PERSON_ID"],
    }
    KNOWN_FIELDS = [
        "PERSON_ID",
        "DATE_OF_BIRTH",
        "YEAR_OF_BIRTH",
        "DATE_OF_DEATH",
        "SEX",
        "ETHNICITY",
    ]
    DEFAULT_MAPPING = {"PERSON_ID": "PERSON_ID"}


class EventTable(PhenexTable):
    NAME_TABLE = "EVENT"
    KNOWN_FIELDS = ["PERSON_ID", "EVENT_DATE"]
    DEFAULT_MAPPING = {
        "PERSON_ID": "PERSON_ID",
        "EVENT_DATE": "EVENT_DATE",
    }


class CodeTable(PhenexTable):
    """
    Base class for tables containing coded events (diagnoses, procedures, medications, etc.).

    CODES_DEFINED_IN: Specifies where CODE/CODE_TYPE columns are located.
        - None (default): Codes are in this table itself
        - NAME_TABLE or class name: Codes are in another table (matched by NAME_TABLE or class name)

    Example 1: Traditional pattern (codes in the table)
        class ConditionOccurrenceTable(CodeTable):
            CODES_DEFINED_IN = None  # Default, codes are here
            DEFAULT_MAPPING = {
                "CODE": "CONDITION_CONCEPT_ID",
                "CODE_TYPE": "VOCABULARY_ID"
            }

    Example 2: Concept table pattern (codes in separate table)
        class EventTable(CodeTable):
            CODES_DEFINED_IN = "CONCEPT"  # NAME_TABLE of target table
            JOIN_KEYS = {"EventMappingTable": ["EVENTMAPPINGID"]}
            PATHS = {"ConceptTable": ["EventMappingTable"]}
    """

    NAME_TABLE = "CODE"
    RELATIONSHIPS = {
        "PhenexPersonTable": ["PERSON_ID"],
        "PhenexVisitOccurrenceTable": ["PERSON_ID", "VISIT_DETAIL_ID"],
    }
    CODES_DEFINED_IN = None  # Set to domain name if codes are in a different table
    KNOWN_FIELDS = ["PERSON_ID", "EVENT_DATE", "CODE", "CODE_TYPE", "VISIT_DETAIL_ID"]
    DEFAULT_MAPPING = {
        "PERSON_ID": "PERSON_ID",
        "EVENT_DATE": "EVENT_DATE",
        "CODE": "CODE",
    }


class PhenexVisitOccurrenceTable(PhenexTable):
    NAME_TABLE = "VISIT_OCCURRENCE"
    RELATIONSHIPS = {
        "PhenexPersonTable": ["PERSON_ID"],
        "CodeTable": ["PERSON_ID", "VISIT_OCCURRENCE_ID"],
    }
    KNOWN_FIELDS = [
        "PERSON_ID",
        "VISIT_OCCURRENCE_ID",
        "VISIT_OCCURRENCE_SOURCE_VALUE",
    ]

    DEFAULT_MAPPING = {
        "PERSON_ID": "PERSON_ID",
        "VISIT_OCCURRENCE_ID": "VISIT_DETAIL_ID",
        "VISIT_OCCURRENCE_SOURCE_VALUE": "VISIT_DETAIL_SOURCE_VALUE",
    }


class PhenexIndexTable(PhenexTable):
    NAME_TABLE = "INDEX"
    JOIN_KEYS = {
        "CodeTable": ["PERSON_ID"],
        "PhenexVisitOccurrenceTable": ["PERSON_ID"],
    }
    KNOWN_FIELDS = [
        "PERSON_ID",
        "INDEX_DATE",
    ]
    DEFAULT_MAPPING = {"PERSON_ID": "PERSON_ID", "INDEX_DATE": "INDEX_DATE"}


class PhenexObservationPeriodTable(PhenexTable):
    NAME_TABLE = "OBSERVATION_PERIOD"
    KNOWN_FIELDS = [
        "PERSON_ID",
        "START_DATE",
        "END_DATE",
    ]
    DEFAULT_MAPPING = {
        "PERSON_ID": "PERSON_ID",
        "START_DATE": "OBSERVATION_PERIOD_START_DATE",
        "END_DATE": "OBSERVATION_PERIOD_END_DATE",
    }


class MeasurementTable(PhenexTable):
    NAME_TABLE = "MEASUREMENT"
    RELATIONSHIPS = {
        "PhenexPersonTable": ["PERSON_ID"],
        "PhenexVisitOccurrenceTable": ["PERSON_ID", "VISIT_DETAIL_ID"],
    }
    KNOWN_FIELDS = [
        "PERSON_ID",
        "EVENT_DATE",
        "CODE",
        "CODE_TYPE",
        "VISIT_DETAIL_ID",
        "VALUE",
    ]
    DEFAULT_MAPPING = {
        "PERSON_ID": "PERSON_ID",
        "EVENT_DATE": "EVENT_DATE",
        "CODE": "CODE",
        "VALUE": "VALUE",
    }


class PhenotypeTable(PhenexTable):
    NAME_TABLE = "PHENOTYPE"
    KNOWN_FIELDS = ["PERSON_ID", "BOOLEAN", "EVENT_DATE", "VALUE"]
    DEFAULT_MAPPING = {
        "PERSON_ID": "PERSON_ID",
        "BOOLEAN": "BOOLEAN",
        "EVENT_DATE": "EVENT_DATE",
        "VALUE": "VALUE",
    }


def is_phenex_person_table(table: PhenexTable) -> bool:
    """
    Check if given table is a person table.
    One could check one row per patient?
    """
    return set(table.columns) >= set(PhenexPersonTable.REQUIRED_FIELDS)


def is_phenex_code_table(table: PhenexTable) -> bool:
    """
    Check if given table is a code table.
    """
    return set(table.columns) >= set(CodeTable.REQUIRED_FIELDS)


def is_phenex_event_table(table: PhenexTable) -> bool:
    """
    Check if given table is a code table.
    """
    return set(table.columns) >= set(EventTable.REQUIRED_FIELDS)


def is_phenex_phenotype_table(table: PhenexTable) -> bool:
    """
    Check if given table is a code table.
    """
    return set(table.columns) >= set(PhenotypeTable.REQUIRED_FIELDS)


def is_phenex_index_table(table: PhenexTable) -> bool:
    """
    Check if given table is a code table.
    """
    return set(table.columns) >= set(PhenexIndexTable.REQUIRED_FIELDS)


PHENOTYPE_TABLE_COLUMNS = ["PERSON_ID", "BOOLEAN", "EVENT_DATE", "VALUE"]
