import sqlalchemy as sa
from sqlalchemy.orm import Session

from openlog.datamodel.assay.generic_assay import (
    AssayColumn,
    AssayDatabaseDefinition,
    AssayDataExtent,
    AssayDefinition,
    AssayDomainType,
    AssaySeriesType,
    DepthAssay,
    GenericAssay,
    TimeAssay,
)
from openlog.datamodel.connection.sqlalchemy.assay_reader import (
    SqlAlchemyDiscreteDepthAssay,
    SqlAlchemyDiscreteTimeAssay,
    SqlAlchemyExtendedDepthAssay,
    SqlAlchemyExtendedTimeAssay,
)


class AssayFactory:
    def __init__(
        self, assay_definition: AssayDefinition, assay_table: AssayDatabaseDefinition
    ):
        self._assay_definition = assay_definition
        self._assay_table = assay_table

    def create_assay_table(self) -> sa.Table:
        """
        Create sqlalchemy table for assay
        Returns: sqlalchemy table

        """
        metadata_obj = sa.MetaData()

        table = sa.Table(
            self._assay_table.table_name,
            metadata_obj,
            sa.Column(self._assay_table.hole_id_col, sa.String, primary_key=True),
            self.create_x_column(self._assay_table.x_col, True),
            sa.Column(self._assay_table.dataset_col, sa.String),
        )

        # Add params for each column
        for col, col_def in self._assay_definition.columns.items():
            database_column = self._assay_table.y_col[col]
            table.append_column(
                self.create_y_column(col_def, database_column), replace_existing=False
            )

            # Add image format column if defined
            if col_def.image_format_col:
                table.append_column(
                    sa.Column(col_def.image_format_col, sa.String()),
                    replace_existing=False,
                )

        # Add column uncertainty
        for col in self._assay_definition.get_uncertainty_columns():
            table.append_column(sa.Column(col, sa.Float), replace_existing=False)

        if self._assay_table.schema:
            table.schema = self._assay_table.schema

        if self._assay_definition.data_extent == AssayDataExtent.EXTENDED:
            table.append_column(
                self.create_x_column(self._assay_table.x_end_col, True),
                replace_existing=False,
            )
        return table

    def create_x_column(self, x_col: str, as_primary_key: bool = False) -> sa.Column:
        """
        Create a sqlalchemy column for x value definition depending on assay domain

        Args:
            x_col: x column name
            as_primary_key: column used as primary key

        Returns: sqlalchemy column

        """
        col_type = sa.Numeric(asdecimal=False)
        if self._assay_definition.domain == AssayDomainType.TIME:
            col_type = sa.DateTime()

        return sa.Column(x_col, col_type, primary_key=as_primary_key)

    @staticmethod
    def create_y_column(assay_column: AssayColumn, y_col: str) -> sa.Column:
        """
        Create a sqlalchemy column for y value definition depending on assay series type

        Args:
            assay_column: AssayColumn
            y_col : y column name

        Returns: sqlalchemy column

        """
        # Add support of CATEGORICAL for lith assay display, not available yet in assay import
        if assay_column.series_type == AssaySeriesType.NUMERICAL:
            col_type = sa.Float()
        elif assay_column.series_type == AssaySeriesType.CATEGORICAL:
            col_type = sa.String()
        elif assay_column.series_type == AssaySeriesType.NOMINAL:
            col_type = sa.String()
        elif assay_column.series_type == AssaySeriesType.DATETIME:
            col_type = sa.DateTime()
        elif assay_column.series_type == AssaySeriesType.IMAGERY:
            col_type = sa.LargeBinary()
        else:
            col_type = sa.Float()

        return sa.Column(y_col, col_type)

    def create_generic_assay(
        self,
        hole_id: str,
        session: Session,
    ) -> GenericAssay:
        """
        Create a GenericAssay for sqlalchemy session and base use depending on AssayDefinition

        Args:
            hole_id: collar hole_id
            session: sqlalchemy session

        Returns: GenericAssay for AssayDefinition

        """
        if self._assay_definition.domain == AssayDomainType.DEPTH:
            return self._create_depth_assay(hole_id, session)
        else:
            return self._create_time_assay(hole_id, session)

    def _create_depth_assay(
        self,
        hole_id: str,
        session: Session,
    ) -> DepthAssay:
        """
        Create a DepthAssay for sqlalchemy session and base use depending on AssayDefinition

        Args:
            hole_id: collar hole_id
            session: sqlalchemy session

        Returns: DepthAssay for AssayDefinition

        """
        assay_table = self.create_assay_table()
        if self._assay_definition.data_extent == AssayDataExtent.DISCRETE:
            return SqlAlchemyDiscreteDepthAssay(
                hole_id, self._assay_definition, self._assay_table, session, assay_table
            )
        else:
            return SqlAlchemyExtendedDepthAssay(
                hole_id, self._assay_definition, self._assay_table, session, assay_table
            )

    def _create_time_assay(
        self,
        hole_id: str,
        session: Session,
    ) -> TimeAssay:
        """
        Create a TimeAssay for sqlalchemy session and base use depending on AssayDefinition

        Args:
            hole_id: collar hole_id
            session: sqlalchemy session

        Returns: TimeAssay for AssayDefinition

        """
        assay_table = self.create_assay_table()
        if self._assay_definition.data_extent == AssayDataExtent.DISCRETE:
            return SqlAlchemyDiscreteTimeAssay(
                hole_id, self._assay_definition, self._assay_table, session, assay_table
            )
        else:
            return SqlAlchemyExtendedTimeAssay(
                hole_id, self._assay_definition, self._assay_table, session, assay_table
            )
