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

from unittest.mock import AsyncMock, MagicMock, patch
from io import BytesIO

import pandas as pd
import pytest
import pytest_asyncio
from sqlalchemy import create_engine, text
from testcontainers.minio import MinioContainer
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 minio_container():
    """
    MinIO S3-compatible storage for E2E tests that exercise real offload uploads.
    """
    minio = MinioContainer()
    minio.start()
    yield minio
    minio.stop()


@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.fixture
def notification_inserts(mock_mongo_client):
    """
    Mongo insert_one mock used by CoreNotificationHandler for notification persistence.

    Yields:
        MagicMock for insert_one, reset before each test.
    """
    mock_db = mock_mongo_client.__getitem__.return_value
    mock_collection = mock_db.__getitem__.return_value
    mock_collection.insert_one.reset_mock()
    yield mock_collection.insert_one


@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()

    # Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve()
    # can decode the payload if offloading is exercised in an integration scenario.
    parquet_df = pd.DataFrame({'a': [1]})
    parquet_buffer = BytesIO()
    parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
    parquet_bytes = parquet_buffer.getvalue()
    
    # sientia_do MinioRepository API
    mock_repo.bucket = 'test-bucket'
    mock_repo.upload_file = AsyncMock(
        side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: {
            'minio_object_name': f'sientia/streamlit-connectors/{relative_key}',
            'original_filename': relative_key.rsplit('/', 1)[-1],
            'uploaded_at': '2024-01-01T00:00:00Z',
            'sha256_hash': 'deadbeef',
        }
    )
    mock_repo.download_file = AsyncMock(return_value=parquet_bytes)
    mock_repo.list_objects = AsyncMock(return_value=[])
    mock_repo.delete_file = AsyncMock()
    mock_repo.close = MagicMock()
    
    return mock_repo


@pytest_asyncio.fixture
def mock_pi_web_api_repository():
    """Mock PI Web API repository for PI Web API operations."""
    mock_repo = MagicMock()

    async def _write_value(web_ids, value, metadata=None, **kwargs):
        """
        Mirror successful PI writes: one response item per requested web_id.

        write_pi_web_api_data passes the list into process_pi_web_api_response (not a
        wrapped {'Items': ...} envelope).
        """
        return [{'WebId': wid, 'Errors': []} for wid in web_ids]

    mock_repo.write_value = AsyncMock(side_effect=_write_value)
    mock_repo.close = MagicMock()
    return mock_repo

@pytest_asyncio.fixture
def mock_opc_repository():
    """Mock OPC repository for OPC operations."""
    mock_repo = MagicMock()
    mock_repo.write_data = AsyncMock(
        return_value=(True, {'response_time': 0.1})
    )
    mock_repo.disconnect = AsyncMock()
    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."""
    # Patch where Activities resolves the symbol (import binds the original class).
    with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository):
        yield

@pytest_asyncio.fixture
def patch_pi_web_api_repository(mock_pi_web_api_repository):
    """Patch MLflowRepository to return mock."""
    with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_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,
    patch_pi_web_api_repository,
    mock_opc_repository
):
    """
    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={
            # Host:port only; Minio() prepends http(s):// from the secure flag.
            'endpoint_url': 'localhost:9000',
            'access_key': 'test',
            'secret_key': 'test',
            'default_bucket': 'test-bucket',
            'retention_hours': 24,
            'secure': False,
        },
        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,
    )

    activities.opc_repository = {
        '1': mock_opc_repository,
    }

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


@pytest_asyncio.fixture(scope='function')
async def test_activities_real_minio(
    postgres_engine,
    postgres_container,
    minio_container,
    mock_logger,
    notification_handler,
    metrics_controller,
    patch_create_engine,
    patch_mlflow,
    patch_pi_web_api_repository,
    mock_opc_repository,
):
    """
    Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests.
    """
    minio_client = minio_container.get_client()
    if not minio_client.bucket_exists('test-bucket'):
        minio_client.make_bucket('test-bucket')
    minio_port = minio_container.get_exposed_port(9000)
    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': f'localhost:{minio_port}',
            'access_key': 'minioadmin',
            'secret_key': 'minioadmin',
            'default_bucket': 'test-bucket',
            'retention_hours': 24,
            'secure': False,
        },
        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,
    )
    activities.opc_repository = {'1': mock_opc_repository}
    try:
        yield activities
    finally:
        await activities.shutdown()


def _worker_activity_list(test_activities: Activities):
    return [
        test_activities.load_custom_query,
        test_activities.load_query_with_minio_offload,
        test_activities.cleanup_minio_objects_expired,
        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.export_payload_to_postgres,
        test_activities.write_metrics,
    ]


@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=_worker_activity_list(test_activities),
    ) as worker:
        yield worker


@pytest_asyncio.fixture(scope='function')
async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio):
    """Temporal worker backed by Activities using real MinIO testcontainer."""
    async with Worker(
        temporal_test_env.client,
        task_queue='test-queue',
        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
        activities=_worker_activity_list(test_activities_real_minio),
    ) as worker:
        yield worker