"""
Pytest configuration and fixtures for E2E tests.
"""

from unittest.mock import AsyncMock, MagicMock, patch

import pandas as pd
import pytest
import pytest_asyncio
from sqlalchemy import create_engine, text
from testcontainers.postgres import PostgresContainer
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController

# Test constants
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'


@pytest_asyncio.fixture(scope='session')
def postgres_container():
    """
    Create a PostgreSQL container using testcontainers.
    
    This fixture creates a real PostgreSQL database in a Docker container
    that will be used for all tests in the session.
    """
    postgres = PostgresContainer('postgres:15')
    postgres.start()
    yield postgres
    postgres.stop()


@pytest_asyncio.fixture
def postgres_engine(postgres_container):
    """
    Create SQLAlchemy engine for PostgreSQL test database.
    
    This fixture creates a connection to the PostgreSQL container
    created by the postgres_container fixture.
    """
    engine = create_engine(postgres_container.get_connection_url())
    
    yield engine
    
    engine.dispose()


def _create_schema_and_tables(engine):
    """
    Helper function to create schema and tables in the given engine.
    
    Creates predictions_schema with:
    - laborious_data: Input data table for queries
    - predictions: Output predictions table
    - transformed_data: Output transformed data table
    """
    # Use begin() to ensure transaction is properly committed
    with engine.begin() as conn:
        # Create predictions_schema
        conn.execute(text("CREATE SCHEMA IF NOT EXISTS predictions_schema"))
        
        # Create laborious_data table (input data from sensors)
        create_laborious_data_sql = """
            CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
                id SERIAL NOT NULL,
                model_id int4 NOT NULL,
                variable text NOT NULL,
                value numeric NULL,
                "timestamp" timestamptz NOT NULL,
                created_at timestamptz NOT NULL,
                PRIMARY KEY (id)
            );
        """
        conn.execute(text(create_laborious_data_sql))
        
        # Create predictions table
        create_predictions_sql = """
            CREATE TABLE if not exists predictions_schema.predictions (
                id SERIAL NOT NULL ,
                model_id int4 NOT NULL,
                prediction numeric NULL,
                prediction_confidence numeric NOT NULL,
                response_time numeric NOT NULL,
                prediction_status text NOT NULL,
                "timestamp" timestamptz NOT NULL,
                created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
                "comments" text NULL,
                PRIMARY KEY (id, created_at)
            );
        """
        conn.execute(text(create_predictions_sql))
        
        # Create transformed_data table
        create_transformed_sql = """
            CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
                id SERIAL NOT NULL,
                model_id int4 NOT NULL,
                variable text NOT NULL,
                value numeric NULL,
                "timestamp" timestamptz NOT NULL,
                created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
                PRIMARY KEY (id)
            );
        """
        conn.execute(text(create_transformed_sql))


@pytest_asyncio.fixture(autouse=True)
def setup_postgres_schema_and_tables(postgres_engine):
    """
    Automatically create necessary schema and tables before each test.
    
    This fixture runs automatically (autouse=True) and ensures
    that the predictions_schema and tables exist with the correct structure.
    """
    _create_schema_and_tables(postgres_engine)
    yield


@pytest_asyncio.fixture
def mock_logger():
    """Mock logger for testing."""
    def message(message):
        print(f"[LOG] {message}")
    def custom_message(message, _metadata={}):
        print(f"[LOG] {message}")
    logger = MagicMock()
    logger.info = MagicMock(
        side_effect=message
    )
    logger.debug = MagicMock(
        side_effect=message
    )
    logger.error = MagicMock(
        side_effect=message
    )
    logger.warning = MagicMock(
        side_effect=message
    )
    logger.custom_info = MagicMock(
        side_effect=custom_message
    )
    logger.custom_debug = MagicMock(
        side_effect=custom_message
    )
    logger.custom_error = MagicMock(
        side_effect=custom_message
    )
    logger.custom_warning = MagicMock(
        side_effect=custom_message
    )
    return logger


@pytest_asyncio.fixture
def mock_mongo_client():
    """
    Mock MongoDB client to avoid real connections.
    
    This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
    allowing us to use a real NotificationHandler instance without connecting to MongoDB.
    """
    mock_client = MagicMock()
    mock_db = MagicMock()
    mock_collection = MagicMock()
    
    # Configure the mock chain: client[database] -> db[collection] -> collection
    mock_client.__getitem__.return_value = mock_db
    mock_db.__getitem__.return_value = mock_collection
    
    # Mock server_info() to avoid connection attempts
    mock_client.server_info = MagicMock()
    
    # Mock insert_one for notifications
    mock_collection.insert_one = MagicMock()
    
    return mock_client


@pytest_asyncio.fixture
def notification_handler(mock_logger, mock_mongo_client):
    """
    Create a real NotificationHandler instance with mocked MongoDB client.
    
    This fixture creates a real CoreNotificationHandler instance but mocks
    the underlying MongoDB connection to avoid real database connections.
    """
    # Patch MongoClient where it's imported in the handlers module
    with patch('sientia_do.notifications.handlers.MongoClient', return_value=mock_mongo_client):
        handler = CoreNotificationHandler(
            connection_string=TEST_MONGODB_CONNECTION_STRING,
            database=TEST_DATABASE_NAME,
            logger=mock_logger,
            project_name='laborious',
        )
        yield handler
        handler.shutdown()


@pytest_asyncio.fixture
def metrics_controller(mock_logger):
    """Create a real MetricsController instance."""
    return MetricsController(logger=mock_logger)


@pytest_asyncio.fixture
def mock_minio_repository():
    """Mock MinIO repository for object storage operations."""
    mock_repo = MagicMock()
    
    # Mock repository methods
    mock_repo.put_parquet_from_dataframe = AsyncMock(return_value='test-object-key')
    mock_repo.get_parquet_as_dataframe = AsyncMock(return_value=pd.DataFrame())
    mock_repo.minio_bucket = 'test-bucket'
    
    return mock_repo


@pytest_asyncio.fixture
def patch_create_engine(postgres_engine):
    """Patch create_engine to return test postgres_engine."""
    with patch('sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine):
        yield


@pytest_asyncio.fixture
def patch_minio_repository(mock_minio_repository):
    """Patch MinioRepository to return mock."""
    with patch('laborious.utils.repository.minio_repository.MinioRepository', return_value=mock_minio_repository):
        yield


@pytest_asyncio.fixture
def mock_mlflow_models():
    """Create mock models for MLflow load_model methods."""
    # Mock transform model - returns DataFrame with same index as input
    mock_transform_model = MagicMock()
    def mock_transform_predict(data):
        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
        print(data.to_csv())
        print(data.index)
        result = pd.DataFrame({
            'feature_1': [0.234] * num_rows,
            'feature_2': [0.783] * num_rows,
        })
        result.index = data.index
        return result
    mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
    
    # Mock predict model - returns array/list of predictions
    mock_predict_model = MagicMock()
    def mock_predict_predict(data):
        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
        return [0.5] * num_rows
    mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
    
    # Mock PyFuncModel for compressed models
    mock_pyfunc_model = MagicMock()
    mock_pyfunc_model._model_impl = MagicMock()
    mock_pyfunc_model._model_impl.python_model = mock_transform_model
    
    return {
        'transform_model': mock_transform_model,
        'predict_model': mock_predict_model,
        'pyfunc_model': mock_pyfunc_model,
    }


@pytest_asyncio.fixture
def patch_mlflow(mock_mlflow_models):
    """Patch mlflow module in repository with load_model mocks."""
    mock_mlflow = MagicMock()
    
    # Mock sklearn.load_model
    def mock_sklearn_load_model(model_uri):
        if 'data_model' in model_uri or 'transform' in model_uri.lower():
            return mock_mlflow_models['transform_model']
        return mock_mlflow_models['predict_model']
    mock_mlflow.sklearn = MagicMock()
    mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
    
    # Mock pyfunc.load_model
    def mock_pyfunc_load_model(model_uri):
        if 'artifacts' in model_uri or 'tmp' in model_uri:
            return mock_mlflow_models['pyfunc_model']
        if 'data_model' in model_uri or 'transform' in model_uri.lower():
            return mock_mlflow_models['transform_model']
        return mock_mlflow_models['predict_model']
    mock_mlflow.pyfunc = MagicMock()
    mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
    
    # Mock pytorch.load_model
    mock_mlflow.pytorch = MagicMock()
    mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
    
    # Mock other mlflow methods that might be called
    mock_mlflow.set_tracking_uri = MagicMock()
    mock_mlflow.get_run = MagicMock(return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id')))
    mock_mlflow.tracking = MagicMock()
    mock_mlflow.tracking.MlflowClient = MagicMock(return_value=MagicMock(
        search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
        search_model_versions=MagicMock(return_value=[MagicMock(
            current_stage='Production',
            version='1',
            source='runs:/artifacts/test_run_id'
        )])
    ))
    
    with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
        yield mock_mlflow


@pytest_asyncio.fixture(scope='function')
async def test_activities(
    postgres_engine,
    postgres_container,
    mock_logger,
    notification_handler,
    metrics_controller,
    mock_minio_repository,
    patch_create_engine,
    patch_minio_repository,
    patch_mlflow,
):
    """
    Create Activities instance with test dependencies.

    This fixture creates a real Activities instance with:
    - PostgreSQL database (via testcontainers)
    - Mocked MinIO client
    - Real NotificationHandler and MetricsController (with mocked underlying services)
    """
    activities = Activities(
        postgres_config={
            'host': 'localhost',
            'port': postgres_container.get_exposed_port(5432),
            'user': 'test',
            'password': 'test',
            'dbname': 'test',
            'min_connections': 1,
            'max_connections': 5,
        },
        mlflow_config={
            'host': 'http://localhost',
            'port': '5000',
            'username': 'test',
            'password': 'test',
        },
        minio_config={
            'endpoint_url': 'http://localhost:9000',
            'access_key': 'test',
            'secret_key': 'test',
            'region_name': 'us-east-1',
            'default_bucket': 'test-bucket',
        },
        opc_config={},
        pi_web_api_config={
            'base_url': 'http://localhost:8080',
            'auth_type': 'bearer',
            'auth_token': 'test_token',
        },
        logger=mock_logger,
        notification_handler=notification_handler,
    )

    try:
        yield activities
    finally:
        # Cleanup - ALWAYS runs, even if test fails
        await activities.shutdown()


@pytest_asyncio.fixture(scope='function')
async def temporal_test_env():
    """Create Temporal test environment."""
    env = await WorkflowEnvironment.start_time_skipping()
    async with env:
        yield env


@pytest_asyncio.fixture(scope='function')
async def temporal_worker(temporal_test_env, test_activities):
    """Create Temporal worker with test activities."""
    async with Worker(
        temporal_test_env.client,
        task_queue='test-queue',
        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
        activities=[
            test_activities.load_custom_query,
            test_activities.get_last_timestamp,
            test_activities.input_gate,
            test_activities.request_transform,
            test_activities.mlflow_response_gate,
            test_activities.mlflow_content_gate,
            test_activities.request_predict,
            test_activities.repeat_last_prediction,
            test_activities.format_prediction,
            test_activities.format_transformed_data,
            test_activities.format_default_prediction,
            test_activities.write_pi_web_api_data,
            test_activities.write_opc_data,
            test_activities.export_data_to_postgres,
            test_activities.write_metrics,
        ],
    ) as worker:
        yield worker