# Copyright 2025 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
from typing import Optional

import pandas as pd

import mlrun
import mlrun.common.model_monitoring.helpers
import mlrun.common.schemas.model_monitoring as mm_schemas
import mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_schema as timescaledb_schema
from mlrun.config import config
from mlrun.datastore.datastore_profile import DatastoreProfilePostgreSQL
from mlrun.model_monitoring.db import TSDBConnector
from mlrun.model_monitoring.db.tsdb.preaggregate import (
    PreAggregateConfig,
    PreAggregateManager,
)
from mlrun.model_monitoring.db.tsdb.timescaledb.queries.timescaledb_metrics_queries import (
    TimescaleDBMetricsQueries,
)
from mlrun.model_monitoring.db.tsdb.timescaledb.queries.timescaledb_predictions_queries import (
    TimescaleDBPredictionsQueries,
)
from mlrun.model_monitoring.db.tsdb.timescaledb.queries.timescaledb_results_queries import (
    TimescaleDBResultsQueries,
)
from mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_connection import (
    Statement,
    TimescaleDBConnection,
)
from mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_operations import (
    TimescaleDBOperationsManager,
)
from mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_stream import (
    TimescaleDBStreamProcessor,
)
from mlrun.utils import logger


class TimescaleDBConnector(TSDBConnector):
    """
    Complete TimescaleDB TSDB connector using composition pattern.

    Uses composition for all specialized functionality:
    - TimescaleDBMetricsQueries, TimescaleDBPredictionsQueries, TimescaleDBResultsQueries: Direct query operations
    - TimescaleDBOperationsManager: Table management and write operations
    - TimescaleDBStreamProcessor: Stream processing operations

    Database naming (controlled by mlrun.mlconf.model_endpoint_monitoring.tsdb.auto_create_database):
    - When auto_create_database=True (default): generates database name using system_id: 'mlrun_mm_{system_id}'
    - When auto_create_database=False: uses the database from the profile/connection string as-is
    """

    type: str = mm_schemas.TSDBTarget.TimescaleDB

    def __init__(
        self,
        project: str,
        profile: DatastoreProfilePostgreSQL,
        pre_aggregate_config: Optional[PreAggregateConfig] = None,
        **kwargs,
    ):
        super().__init__(project=project)

        self.profile = profile

        # Determine the monitoring database name
        self.database = self._determine_database_name(profile)

        # Update profile to use the determined database name
        # This ensures the connection uses the correct database
        if profile.database != self.database:
            logger.info(
                "Auto-generated database name for TimescaleDB",
                original_database=profile.database,
                database=self.database,
            )
            # Create a new profile with the generated database
            profile = DatastoreProfilePostgreSQL(
                name=profile.name,
                user=profile.user,
                password=profile.password,
                host=profile.host,
                port=profile.port,
                database=self.database,
            )
            self.profile = profile

        # Create shared connection
        self._connection = TimescaleDBConnection(
            dsn=profile.dsn(),
            min_connections=kwargs.get("min_connections", 1),
            max_connections=kwargs.get("max_connections", 10),
            max_retries=kwargs.get("max_retries", 3),
            retry_delay=kwargs.get("retry_delay", 1.0),
            autocommit=kwargs.get("autocommit", False),
        )

        # Create shared components needed by query classes
        self._tables = timescaledb_schema.create_table_schemas(project)
        self._pre_aggregate_manager = PreAggregateManager(pre_aggregate_config)

        # Create specialized query handlers with proper initialization
        self._metrics_queries = TimescaleDBMetricsQueries(
            project=project,
            connection=self._connection,
            pre_aggregate_manager=self._pre_aggregate_manager,
            tables=self._tables,
        )
        self._predictions_queries = TimescaleDBPredictionsQueries(
            project=project,
            connection=self._connection,
            pre_aggregate_manager=self._pre_aggregate_manager,
            tables=self._tables,
        )
        self._results_queries = TimescaleDBResultsQueries(
            connection=self._connection,
            project=project,
            pre_aggregate_manager=self._pre_aggregate_manager,
            tables=self._tables,
        )

        # Create operations and stream handlers
        self._operations = TimescaleDBOperationsManager(
            project=project,
            connection=self._connection,
            pre_aggregate_config=pre_aggregate_config,
            profile=profile,
        )

        self._stream = TimescaleDBStreamProcessor(
            project=project, profile=profile, connection=self._connection
        )

        self._pre_aggregate_config = pre_aggregate_config

    def _determine_database_name(self, profile: DatastoreProfilePostgreSQL) -> str:
        """
        Determine the database name to use.

        Delegates to the shared helper function to ensure consistent database naming
        across all TimescaleDB components (connector, stream, storey targets).

        :param profile: The PostgreSQL profile
        :return: The database name to use
        """
        return mlrun.common.model_monitoring.helpers.get_tsdb_database_name(
            profile.database
        )

    # Delegate operations methods
    def create_tables(self, *args, **kwargs) -> None:
        return self._operations.create_tables(*args, **kwargs)

    def write_application_event(self, *args, **kwargs) -> None:
        return self._operations.write_application_event(*args, **kwargs)

    def delete_tsdb_records(self, *args, **kwargs) -> None:
        return self._operations.delete_tsdb_records(*args, **kwargs)

    def delete_tsdb_resources(self, *args, **kwargs) -> None:
        return self._operations.delete_tsdb_resources(*args, **kwargs)

    def delete_application_records(self, *args, **kwargs) -> None:
        return self._operations.delete_application_records(*args, **kwargs)

    def read_metrics_data(
        self,
        *,
        endpoint_id: str,
        start: datetime.datetime,
        end: datetime.datetime,
        metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
        type: str,
        with_result_extra_data: bool = False,
    ):
        """Read metrics or results data from TimescaleDB (cross-cutting coordination)."""

        if type == "metrics":
            df = self._metrics_queries.read_metrics_data_impl(
                endpoint_id=endpoint_id,
                start=start,
                end=end,
                metrics=metrics,
            )
            # Use inherited method to convert DataFrame to domain objects
            return self.df_to_metrics_values(
                df=df, metrics=metrics, project=self.project
            )

        else:  # results
            df = self._results_queries.read_results_data_impl(
                endpoint_id=endpoint_id,
                start=start,
                end=end,
                metrics=metrics,
                with_result_extra_data=with_result_extra_data,
            )
            # Use inherited method to convert DataFrame to domain objects
            return self.df_to_results_values(
                df=df, metrics=metrics, project=self.project
            )

    def get_model_endpoint_real_time_metrics(self, *args, **kwargs):
        return self._metrics_queries.get_model_endpoint_real_time_metrics(
            *args, **kwargs
        )

    def get_metrics_metadata(self, *args, **kwargs):
        return self._metrics_queries.get_metrics_metadata(*args, **kwargs)

    def add_basic_metrics(
        self,
        model_endpoint_objects: list[mlrun.common.schemas.ModelEndpoint],
        metric_list: Optional[list[str]] = None,
    ) -> list[mlrun.common.schemas.ModelEndpoint]:
        """
        Add basic metrics to the model endpoint object using TimescaleDB optimizations.

        :param model_endpoint_objects: A list of `ModelEndpoint` objects that will
                                        be filled with the relevant basic metrics.
        :param metric_list:            List of metrics to include from the time series DB. Defaults to all metrics.

        :return: A list of `ModelEndpointMonitoringMetric` objects.
        """
        uids = [mep.metadata.uid for mep in model_endpoint_objects]

        # Access methods directly from the respective query classes
        # Note: last_request is handled separately due to potential data synchronization issues
        metric_name_to_function = {
            mm_schemas.EventFieldType.ERROR_COUNT: self._results_queries.get_error_count,
            mm_schemas.ModelEndpointSchema.AVG_LATENCY: self._predictions_queries.get_avg_latency,
            mm_schemas.ResultData.RESULT_STATUS: self._results_queries.get_drift_status,
        }

        if metric_list is not None:
            for metric_name in list(metric_name_to_function):
                if metric_name not in metric_list:
                    del metric_name_to_function[metric_name]

        metric_name_to_df = {
            metric_name: function(endpoint_ids=uids)
            for metric_name, function in metric_name_to_function.items()
        }

        def add_metrics(
            mep: mlrun.common.schemas.ModelEndpoint,
            df_dictionary: dict[str, pd.DataFrame],
        ):
            for metric in df_dictionary:
                df = df_dictionary.get(metric, pd.DataFrame())
                if not df.empty:
                    line = df[
                        df[mm_schemas.WriterEvent.ENDPOINT_ID] == mep.metadata.uid
                    ]
                    if not line.empty and metric in line:
                        value = line[metric].item()
                        if isinstance(value, pd.Timestamp):
                            value = value.to_pydatetime()
                        setattr(mep.status, metric, value)

            return mep

        enriched_endpoints = list(
            map(
                lambda mep: add_metrics(
                    mep=mep,
                    df_dictionary=metric_name_to_df,
                ),
                model_endpoint_objects,
            )
        )

        # Handle last_request separately with special enrichment
        if metric_list is None or "last_request" in metric_list:
            self._enrich_mep_with_last_request(
                model_endpoint_objects={
                    mep.metadata.uid: mep for mep in enriched_endpoints
                }
            )

        return enriched_endpoints

    def _enrich_mep_with_last_request(
        self,
        model_endpoint_objects: dict[str, mlrun.common.schemas.ModelEndpoint],
    ):
        """
        Enrich model endpoint objects with last_request data from predictions table.
        This method handles the special case of last_request which may have timing issues.
        """
        try:
            last_request_df = self._predictions_queries.get_last_request(
                endpoint_ids=list(model_endpoint_objects.keys())
            )

            if not last_request_df.empty:
                for _, row in last_request_df.iterrows():
                    endpoint_id = row.get(mm_schemas.WriterEvent.ENDPOINT_ID)
                    last_request = row.get("last_request")

                    if (
                        endpoint_id in model_endpoint_objects
                        and last_request is not None
                    ):
                        if isinstance(last_request, pd.Timestamp):
                            last_request = last_request.to_pydatetime()
                        model_endpoint_objects[
                            endpoint_id
                        ].status.last_request = last_request
        except Exception as e:
            # Log but don't fail - last_request is not critical for basic functionality
            logger.warning(
                "Failed to enrich model endpoints with last_request data",
                error=mlrun.errors.err_to_str(e),
                endpoint_count=len(model_endpoint_objects),
            )

    def read_predictions(self, *args, **kwargs):
        return self._predictions_queries.read_predictions(*args, **kwargs)

    def _get_records(
        self,
        table: str,
        start: datetime.datetime,
        end: datetime.datetime,
        endpoint_id: Optional[str] = None,
        columns: Optional[list[str]] = None,
        timestamp_column: Optional[str] = None,
    ) -> pd.DataFrame:
        """
        Get raw records from TimescaleDB as pandas DataFrame.

        This method provides direct access to raw table data.

        :param table: Table name - use TimescaleDBTables enum (METRICS, APP_RESULTS, or PREDICTIONS)
        :param start: Start time for the query
        :param end: End time for the query
        :param endpoint_id: Optional endpoint ID filter (None = all endpoints)
        :param columns: Optional list of specific columns to return (None = all columns)
        :param timestamp_column: Optional timestamp column to use for time filtering (None = use table's default)
        :return: Raw pandas DataFrame with all matching records
        """
        if table == mm_schemas.TimescaleDBTables.METRICS:
            df = self._metrics_queries.read_metrics_data_impl(
                endpoint_id=endpoint_id,
                start=start,
                end=end,
                metrics=None,  # Get all metrics
                timestamp_column=timestamp_column,
            )
        elif table == mm_schemas.TimescaleDBTables.APP_RESULTS:
            df = self._results_queries.read_results_data_impl(
                endpoint_id=endpoint_id,
                start=start,
                end=end,
                metrics=None,  # Get all results
                with_result_extra_data=True,
                timestamp_column=timestamp_column,
            )
        elif table == mm_schemas.TimescaleDBTables.PREDICTIONS:
            df = self._predictions_queries.read_predictions_impl(
                endpoint_id=endpoint_id,
                start=start,
                end=end,
                columns=columns,
                timestamp_column=timestamp_column,
            )
        else:
            raise mlrun.errors.MLRunInvalidArgumentError(
                f"Invalid table '{table}'. Must be METRICS, APP_RESULTS, or PREDICTIONS from TimescaleDBTables enum"
            )

        if columns is not None and not df.empty:
            # Filter to requested columns if specified
            available_columns = [col for col in columns if col in df.columns]
            df = df[available_columns]

        return df

    def get_last_request(self, *args, **kwargs):
        return self._predictions_queries.get_last_request(*args, **kwargs)

    def get_avg_latency(self, *args, **kwargs):
        return self._predictions_queries.get_avg_latency(*args, **kwargs)

    def count_processed_model_endpoints(
        self,
        start: Optional[datetime.datetime] = None,
        end: Optional[datetime.datetime] = None,
        application_names: Optional[list[str] | str] = None,
    ) -> dict[str, int]:
        """
        Count unique endpoints per application from METRICS and APP_RESULTS tables.

        Uses SQL UNION to efficiently count endpoints that have data in EITHER table.

        :param start: Start time for the query (default: last 24 hours)
        :param end: End time for the query (default: current time)
        :param application_names: Filter by specific application names
        :return: Dictionary mapping application_name to endpoint count
        """
        # Set default time range
        start = start or (mlrun.utils.datetime_now() - datetime.timedelta(hours=24))
        start, end = self._pre_aggregate_manager.get_start_end(start, end)

        metrics_table = self._tables[mm_schemas.TimescaleDBTables.METRICS]
        app_results_table = self._tables[mm_schemas.TimescaleDBTables.APP_RESULTS]
        time_column = mm_schemas.WriterEvent.END_INFER_TIME
        app_column = mm_schemas.WriterEvent.APPLICATION_NAME
        endpoint_column = mm_schemas.WriterEvent.ENDPOINT_ID

        # Build application filter and params
        app_filter_metrics = ""
        app_filter_results = ""

        if application_names:
            if isinstance(application_names, str):
                application_names = [application_names]
            app_names_list = list(application_names)
            app_placeholders = ", ".join(["%s"] * len(app_names_list))
            app_filter_metrics = f"AND {app_column} IN ({app_placeholders})"
            app_filter_results = f"AND {app_column} IN ({app_placeholders})"
            # Params: metrics (start, end, apps), app_results (start, end, apps)
            params = [start, end] + app_names_list + [start, end] + app_names_list
        else:
            params = [start, end, start, end]

        # Use UNION to combine endpoints from both METRICS and APP_RESULTS tables
        query_sql = f"""
        SELECT {app_column}, COUNT(DISTINCT {endpoint_column}) as endpoint_count
        FROM (
            SELECT DISTINCT {app_column}, {endpoint_column}
            FROM {metrics_table.full_name()}
            WHERE {time_column} >= %s AND {time_column} <= %s
            {app_filter_metrics}

            UNION

            SELECT DISTINCT {app_column}, {endpoint_column}
            FROM {app_results_table.full_name()}
            WHERE {time_column} >= %s AND {time_column} <= %s
            {app_filter_results}
        ) combined
        GROUP BY {app_column}
        """

        stmt = Statement(query_sql, params)
        result = self._connection.run(query=stmt)

        if not result or not result.data:
            return {}

        # Convert result to dict: {application_name: count}
        return {row[0]: row[1] for row in result.data}

    def get_drift_status(self, *args, **kwargs):
        return self._results_queries.get_drift_status(*args, **kwargs)

    def get_results_metadata(self, *args, **kwargs):
        return self._results_queries.get_results_metadata(*args, **kwargs)

    def get_error_count(self, *args, **kwargs):
        return self._results_queries.get_error_count(*args, **kwargs)

    def count_results_by_status(self, *args, **kwargs):
        return self._results_queries.count_results_by_status(*args, **kwargs)

    def apply_monitoring_stream_steps(self, *args, **kwargs) -> None:
        return self._stream.apply_monitoring_stream_steps(*args, **kwargs)

    def handle_model_error(self, *args, **kwargs) -> None:
        return self._stream.handle_model_error(*args, **kwargs)

    def calculate_latest_metrics(self, *args, **kwargs):
        return self._metrics_queries.calculate_latest_metrics(*args, **kwargs)

    def get_drift_data(self, *args, **kwargs):
        return self._results_queries.get_drift_data(*args, **kwargs)

    def add_pre_writer_steps(self, graph, after):
        return graph.add_step(
            "mlrun.model_monitoring.db.tsdb.timescaledb.writer_graph_steps.ProcessBeforeTimescaleDBWriter",
            name="ProcessBeforeTimescaleDBWriter",
            after=after,
        )

    def apply_writer_steps(self, graph, after, **kwargs) -> None:
        tables = timescaledb_schema.create_table_schemas(self.project)

        graph.add_step(
            "mlrun.datastore.storeytargets.TimescaleDBStoreyTarget",
            name="tsdb_metrics",
            after=after,
            url=f"ds://{self.profile.name}",
            table=tables[mm_schemas.TimescaleDBTables.METRICS].full_name(),
            time_col=mm_schemas.WriterEvent.END_INFER_TIME,
            columns=[
                mm_schemas.WriterEvent.START_INFER_TIME,
                mm_schemas.MetricData.METRIC_VALUE,
                mm_schemas.WriterEvent.ENDPOINT_ID,
                mm_schemas.WriterEvent.APPLICATION_NAME,
                mm_schemas.MetricData.METRIC_NAME,
            ],
            max_events=config.model_endpoint_monitoring.writer_graph.max_events,
            flush_after_seconds=config.model_endpoint_monitoring.writer_graph.flush_after_seconds,
        )

        graph.add_step(
            "mlrun.datastore.storeytargets.TimescaleDBStoreyTarget",
            name="tsdb_app_results",
            after=after,
            url=f"ds://{self.profile.name}",
            table=tables[mm_schemas.TimescaleDBTables.APP_RESULTS].full_name(),
            time_col=mm_schemas.WriterEvent.END_INFER_TIME,
            columns=[
                mm_schemas.WriterEvent.START_INFER_TIME,
                mm_schemas.ResultData.RESULT_VALUE,
                mm_schemas.ResultData.RESULT_STATUS,
                mm_schemas.ResultData.RESULT_EXTRA_DATA,
                mm_schemas.WriterEvent.ENDPOINT_ID,
                mm_schemas.WriterEvent.APPLICATION_NAME,
                mm_schemas.ResultData.RESULT_NAME,
                mm_schemas.ResultData.RESULT_KIND,
            ],
            max_events=config.model_endpoint_monitoring.writer_graph.max_events,
            flush_after_seconds=config.model_endpoint_monitoring.writer_graph.flush_after_seconds,
        )
