from typing import Dict, List

from geoalchemy2 import Geometry
from sqlalchemy import Column, Integer, String
from sqlalchemy.dialects.postgresql import REGCLASS
from sqlalchemy.orm import declarative_base
from xplordb.datamodel.collar import Collar

from openlog.datamodel.assay.generic_assay import (
    AssayColumn,
    AssayDatabaseDefinition,
    AssayDataExtent,
    AssayDefinition,
    AssayDomainType,
    AssaySeriesType,
)
from openlog.datamodel.assay.uncertainty import AssayColumnUncertainty

Base = declarative_base()

DEFAULT_SRID = 3857
DBGEO_PERSON_CODE: str = "dbgeo"
DBGEO_DATASET: str = "dbgeo_dataset"


class MeasureMetadata(Base):
    """
    Define dbgeo columns for measure metadata, used to define available Assay
    """

    __tablename__ = "measure_metadata"
    __table_args__ = {"schema": "measure"}

    measure_table = Column("measure_table", REGCLASS, primary_key=True)
    name = Column("name", String)
    unit = Column("unit_of_measure", String)
    x_axis_type = Column("x_axis_type", String)
    storage_type = Column("storage_type", String)
    measure_category = Column("measure_category", String)

    @property
    def table_name(self) -> str:
        return f"measure_{self.measure_table.replace('measure.', '')}"

    @property
    def assay_schema(self) -> str:
        return "qgis"

    @property
    def data_extent(self) -> AssayDataExtent:
        result = (
            AssayDataExtent.EXTENDED
            if self.storage_type == "Cumulative" or self.storage_type == "Image"
            else AssayDataExtent.DISCRETE
        )
        return result

    @property
    def domain(self) -> AssayDomainType:
        result = (
            AssayDomainType.DEPTH
            if self.x_axis_type == "DepthAxis"
            else AssayDomainType.TIME
        )
        return result

    def _get_available_cols(self, session) -> [str]:
        return session.execute(
            f"SELECT * FROM {self.assay_schema}.{self.table_name} LIMIT 1"
        ).keys()

    def get_filter_column(self, session) -> str:
        result = ""
        # Check if chemical_element columns available
        available_cols = self._get_available_cols(session)
        if "chemical_element" in available_cols:
            # If activity is contained in measure table, radionuclid column must be used
            if "activity" in self.measure_table and "radionuclide" in available_cols:
                result = "radionuclide"
            else:
                result = "chemical_element"
        return result

    def get_filter_column_values(self, session) -> List[str]:
        values = []

        col = self.get_filter_column(session)

        # Get distinct values
        if col:
            values = [
                res[0]
                for res in session.execute(
                    f"SELECT DISTINCT {col} FROM {self.measure_table} WHERE measure_value != 'NaN' and measure_value IS NOT NULL and {col} IS NOT NULL ORDER BY {col}"
                )
            ]
        return values

    def get_uncertainty_column(self, session) -> str:
        uncertainty_column = ""
        # Check if measure_uncertainty column available
        if "measure_uncertainty" in self._get_available_cols(session):
            uncertainty_column = "measure_uncertainty"
        return uncertainty_column

    def to_assay_definition(self, session) -> AssayDefinition:
        filter_column_values = self.get_filter_column_values(session)
        uncertainty = AssayColumnUncertainty(
            upper_whisker_column=self.get_uncertainty_column(session)
        )

        if self.storage_type == "Image":
            serie_type = AssaySeriesType.IMAGERY
            image_format_col = "image_format"
            unit = "m"
        else:
            serie_type = AssaySeriesType.NUMERICAL
            image_format_col = ""
            unit = self.unit

        if len(filter_column_values) == 0:
            columns = {
                self.name: AssayColumn(
                    name=self.name,
                    series_type=serie_type,
                    unit=unit,
                    uncertainty=uncertainty,
                    stored_in_base_unit=False,
                    image_format_col=image_format_col,
                )
            }
        else:
            columns = {}
            for elem in filter_column_values:
                columns[elem] = AssayColumn(
                    elem,
                    series_type=serie_type,
                    unit=self.unit,
                    uncertainty=uncertainty,
                    stored_in_base_unit=False,
                    image_format_col=image_format_col,
                )

        return AssayDefinition(
            variable=self.measure_table,
            domain=self.domain,
            data_extent=self.data_extent,
            columns=columns,
            display_name=self.name,
        )

    def to_assay_database_definition(self, session) -> AssayDatabaseDefinition:
        hole_id_col = "station_id"
        dataset_col = "dataset_id"

        y_col, y_column_filter = self._get_y_col_and_column_filter(session)
        x_col, x_end_col = self._get_x_col_x_end_col()

        return AssayDatabaseDefinition(
            table_name=self.table_name,
            hole_id_col=hole_id_col,
            dataset_col=dataset_col,
            x_col=x_col,
            y_col=y_col,
            y_column_filter=y_column_filter,
            x_end_col=x_end_col,
            schema=self.assay_schema,
        )

    def _get_x_col_x_end_col(self) -> (str, str):
        domain = self.domain
        if self.data_extent == AssayDataExtent.EXTENDED:
            x_col = (
                "start_measure_time" if domain == AssayDomainType.TIME else "depth_from"
            )
            x_end_col = (
                "end_measure_time" if domain == AssayDomainType.TIME else "depth_to"
            )
        else:
            x_col = (
                "measure_time" if domain == AssayDomainType.TIME else "measure_depth"
            )
            x_end_col = ""
        return x_col, x_end_col

    def _get_y_col_and_column_filter(self, session) -> (Dict[str, str], Dict[str, str]):
        y_column_name = self._get_y_column_name(session)
        chemical_elements = self.get_filter_column_values(session)
        if len(chemical_elements) == 0:
            y_col = {self.name: y_column_name}
            y_column_filter = {}
        else:
            col = self.get_filter_column(session)
            y_col = {}
            y_column_filter = {}
            for elem in chemical_elements:
                y_col[elem] = y_column_name
                y_column_filter[elem] = col
        return y_col, y_column_filter

    def _get_y_column_name(self, session) -> str:
        cols = self._get_available_cols(session)
        if "value" in cols:
            y_column_name = "value"
        elif "measure_value" in cols:
            y_column_name = "measure_value"
        elif "image_data" in cols:
            y_column_name = "image_data"
        else:
            y_column_name = cols[0]
        return y_column_name


class BDGeoCollar(Collar, Base):
    """
    Define dbgeo columns for Collar definition
    """

    __tablename__ = "station"
    __table_args__ = {"schema": "qgis"}

    hole_id = Column("id", String, primary_key=True)
    name = Column("name", String)
    data_set = Column("dataset_id", String)
    srid = Column("orig_srid", Integer)
    loaded_by = DBGEO_PERSON_CODE
    geom = Column(
        "point",
        Geometry(
            geometry_type="POINT", dimension=3, management=True, srid=DEFAULT_SRID
        ),
    )

    @property
    def x(self) -> float:
        return self.geom.ST_X()

    @property
    def y(self) -> float:
        return self.geom.ST_Y()

    @property
    def z(self) -> float:
        return self.geom.ST_Z()
