# Copyright (c) 2025 Scientific Informatics, LLC
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/omopcloudetl_core

from datetime import datetime, timezone
from typing import TYPE_CHECKING, Optional, Union, List, Any
from uuid import UUID

from omopcloudetl_core.models.metrics import ExecutionMetrics, LoadMetrics
from omopcloudetl_core.sql_tools import apply_query_tag

if TYPE_CHECKING:
    from omopcloudetl_core.abstractions.connections import BaseConnection


class MetadataManager:
    """Manages logging of workflow execution metadata to a database table."""

    METADATA_TABLE = "omopcloudetl_execution_log"

    def __init__(self, connection: "BaseConnection"):
        """
        Initializes the MetadataManager.

        Args:
            connection: An instance of a BaseConnection implementation.
        """
        self.connection = connection

    def initialize_store(self, execution_id: UUID) -> None:
        """
        Creates the metadata logging table if it does not exist.
        """
        # This DDL is safe from injection as it's static
        sql = f"""
        CREATE TABLE IF NOT EXISTS {self.METADATA_TABLE} (
            log_id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
            execution_id VARCHAR(36) NOT NULL,
            step_name VARCHAR(255) NOT NULL,
            status VARCHAR(50) NOT NULL,
            start_time TIMESTAMP,
            end_time TIMESTAMP,
            duration_seconds DOUBLE PRECISION,
            rows_affected BIGINT,
            rows_inserted BIGINT,
            rows_updated BIGINT,
            rows_deleted BIGINT,
            rows_processed BIGINT,
            rows_rejected BIGINT,
            error_details_uri VARCHAR(2048),
            query_id VARCHAR(255),
            error_message TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """
        tagged_sql = apply_query_tag(sql, {"omopcloudetl_tool": "MetadataManager", "execution_id": str(execution_id)})
        self.connection.execute_sql(tagged_sql)

    def log_step_start(self, execution_id: UUID, step_name: str) -> None:
        """
        Logs the start of a workflow step using parameterized queries.
        """
        sql = f"""
        INSERT INTO {self.METADATA_TABLE}
            (execution_id, step_name, status, start_time)
        VALUES
            (?, ?, ?, ?);
        """
        params = [str(execution_id), step_name, "RUNNING", datetime.now(timezone.utc)]
        tagged_sql = apply_query_tag(
            sql,
            {
                "omopcloudetl_tool": "MetadataManager",
                "execution_id": str(execution_id),
                "step_name": step_name,
            },
        )
        self.connection.execute_sql(tagged_sql, params=params)

    def log_step_end(
        self,
        execution_id: UUID,
        step_name: str,
        metrics: Union[ExecutionMetrics, LoadMetrics, None],
        error: Optional[Exception] = None,
    ) -> None:
        """
        Logs the end of a workflow step, updating its status and metrics
        using parameterized queries.
        """
        end_time = datetime.now(timezone.utc)
        status = "FAILED" if error else "COMPLETED"

        update_clauses: List[str] = []
        params: List[Any] = []

        update_clauses.append("status = ?")
        params.append(status)

        update_clauses.append("end_time = ?")
        params.append(end_time)

        # Let the database calculate the duration.
        # NOTE: This uses EXTRACT(EPOCH ...), which is part of the SQL standard
        # but may not be supported by all database backends.
        update_clauses.append("duration_seconds = EXTRACT(EPOCH FROM (end_time - start_time))")

        if metrics:
            if isinstance(metrics, ExecutionMetrics):
                if metrics.rows_affected is not None:
                    update_clauses.append("rows_affected = ?")
                    params.append(metrics.rows_affected)
                if metrics.rows_inserted is not None:
                    update_clauses.append("rows_inserted = ?")
                    params.append(metrics.rows_inserted)
            elif isinstance(metrics, LoadMetrics):
                if metrics.rows_processed is not None:
                    update_clauses.append("rows_processed = ?")
                    params.append(metrics.rows_processed)
                if metrics.rows_inserted is not None:
                    update_clauses.append("rows_inserted = ?")
                    params.append(metrics.rows_inserted)

            if metrics.query_id:
                update_clauses.append("query_id = ?")
                params.append(metrics.query_id)

        if error:
            update_clauses.append("error_message = ?")
            params.append(str(error))

        # Final WHERE clause parameters
        params.extend([str(execution_id), step_name, "RUNNING"])

        sql = f"""
        UPDATE {self.METADATA_TABLE}
        SET {", ".join(update_clauses)}
        WHERE
            execution_id = ? AND
            step_name = ? AND
            status = ?;
        """
        tagged_sql = apply_query_tag(
            sql,
            {
                "omopcloudetl_tool": "MetadataManager",
                "execution_id": str(execution_id),
                "step_name": step_name,
            },
        )
        self.connection.execute_sql(tagged_sql, params=params)
