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

import asyncio
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest.mock import MagicMock

import pytest
import pytest_asyncio
from google.protobuf.duration_pb2 import Duration
from pymongo import MongoClient
from redis import Redis
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.temporal.worker.prepare_worker import build_queue_name
from sqlalchemy import create_engine
from temporalio.api.enums.v1 import IndexedValueType
from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest
from temporalio.api.workflowservice.v1 import (
    DescribeNamespaceRequest,
    RegisterNamespaceRequest,
)
from temporalio.client import Client
from temporalio.common import SearchAttributeKey
from temporalio.service import RPCError, RPCStatusCode
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from testcontainers.core.container import DockerContainer
from testcontainers.postgres import PostgresContainer

from e2e.helpers import MONGO_COLLECTIONS, ORCHESTRATOR_TASK_QUEUE
from e2e.smtp_test_server import SmtpTestServer
from e2e.stub_workflows import STUB_WORKFLOW_CLASSES
from orchestrator.activities.activities import Activities
from orchestrator.activities.formatters import schedule_types
from orchestrator.workflows.alerts import Alerts
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.workflows.reports import Reports
from orchestrator.workflows.subworkflows.load_notification_package import (
    LoadNotificationPackage,
)
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications

DB_SCHEMA_SQL_PATH = Path(__file__).parent / 'db_schema.sql'
E2E_DATABASE = 'orchestrator_test'
E2E_RUNTIMES = ('legacy', 'gpu')
MANAGED_NAMESPACES = ('scouter', 'laborious')

E2E_SEARCH_ATTRIBUTES = [
    SearchAttributeKey.for_keyword('model_id'),
    SearchAttributeKey.for_keyword('model_name'),
    SearchAttributeKey.for_keyword('orchestrated'),
]

E2E_NAMESPACE_SEARCH_ATTRIBUTES = {
    'model_id': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
    'model_name': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
    'orchestrated': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
}


async def register_namespace_if_missing(env: WorkflowEnvironment, namespace: str) -> None:
    """
    Register a Temporal namespace on the local dev server and wait until it is ready.

    Args:
        env: Session WorkflowEnvironment from start_local().
        namespace: Namespace name to register.
    """
    service = env.client.service_client
    try:
        await service.workflow_service.register_namespace(
            RegisterNamespaceRequest(
                namespace=namespace,
                workflow_execution_retention_period=Duration(seconds=86400),
            )
        )
    except RPCError as err:
        if err.status != RPCStatusCode.ALREADY_EXISTS:
            raise

    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        try:
            await service.workflow_service.describe_namespace(
                DescribeNamespaceRequest(namespace=namespace)
            )
            return
        except RPCError:
            await asyncio.sleep(0.1)
    raise TimeoutError(f'Namespace {namespace} not ready within 5s')


async def ensure_namespace_search_attributes(
    env: WorkflowEnvironment, namespace: str
) -> None:
    """
    Register the orchestrator search attributes on a namespace, if missing.

    The local Temporal dev server only registers search attributes on the default
    namespace at start time. Schedules created in additional namespaces fail with
    "no mapping defined for search attribute ..." unless we explicitly add the
    same attribute mappings to those namespaces via the operator service.

    Args:
        env: Session WorkflowEnvironment from start_local().
        namespace: Namespace where attributes must be available.
    """
    service = env.client.service_client
    try:
        await service.operator_service.add_search_attributes(
            AddSearchAttributesRequest(
                namespace=namespace,
                search_attributes=dict(E2E_NAMESPACE_SEARCH_ATTRIBUTES),
            )
        )
    except RPCError as err:
        if err.status != RPCStatusCode.ALREADY_EXISTS:
            raise


def temporal_host_from_env(env: WorkflowEnvironment) -> str:
    """Return target host:port for the in-process Temporal dev server."""
    return env.client.service_client.config.target_host


def mongo_uri_from_container(mongo_container) -> str:
    """Build a Mongo connection string for the testcontainer."""
    port = mongo_container.get_exposed_port(27017)
    return f'mongodb://localhost:{port}'


@pytest_asyncio.fixture(scope='session')
def postgres_container():
    """PostgreSQL testcontainer used by all E2E tests."""
    postgres = PostgresContainer('postgres:15')
    postgres.start()
    yield postgres
    postgres.stop()


@pytest_asyncio.fixture(scope='session')
def mongo_container():
    """MongoDB testcontainer used by real CoreNotificationHandler."""
    mongo = DockerContainer('mongo:7').with_exposed_ports(27017)
    mongo.start()
    yield mongo
    mongo.stop()


@pytest_asyncio.fixture(scope='session')
def redis_container():
    """Redis testcontainer for slot and notification timestamp paths."""
    redis = DockerContainer('redis:7').with_exposed_ports(6379)
    redis.start()
    yield redis
    redis.stop()


@pytest_asyncio.fixture
def postgres_engine(postgres_container):
    """SQLAlchemy engine bound to the PostgreSQL testcontainer."""
    engine = create_engine(postgres_container.get_connection_url())
    yield engine
    engine.dispose()


def _create_schema_and_tables(engine):
    sql_text = DB_SCHEMA_SQL_PATH.read_text(encoding='utf-8')
    with engine.begin() as conn:
        conn.exec_driver_sql(sql_text)


@pytest_asyncio.fixture(autouse=True)
def setup_postgres_schema_and_tables(postgres_engine):
    """Recreate Postgres schema from e2e/db_schema.sql before each test."""
    _create_schema_and_tables(postgres_engine)
    yield


@pytest.fixture
def mongo_uri(mongo_container):
    return mongo_uri_from_container(mongo_container)


@pytest_asyncio.fixture(autouse=True)
def reset_mongo_collections(mongo_uri):
    """Drop orchestrator-managed Mongo collections between tests."""
    client = MongoClient(mongo_uri)
    try:
        db = client[E2E_DATABASE]
        for name in MONGO_COLLECTIONS:
            db[name].drop()
    finally:
        client.close()
    yield


@pytest.fixture
def redis_client(redis_container):
    """Redis client bound to the testcontainer."""
    port = int(redis_container.get_exposed_port(6379))
    client = Redis(host='localhost', port=port, decode_responses=True)
    yield client
    client.close()


@pytest_asyncio.fixture(autouse=True)
def reset_redis(redis_client):
    """Flush Redis between tests."""
    redis_client.flushdb()
    yield


@pytest.fixture(scope='session')
def smtp_server():
    """Session-scoped in-process SMTP server."""
    server = SmtpTestServer()
    server.start()
    yield server
    server.stop()


@pytest.fixture(autouse=True)
def smtp_messages_cleanup(smtp_server):
    """Clear captured SMTP messages between tests."""
    smtp_server.clear()
    yield


@pytest_asyncio.fixture(scope='session')
async def temporal_env():
    """Real Temporal dev server (schedule APIs supported)."""
    env = await WorkflowEnvironment.start_local(search_attributes=E2E_SEARCH_ATTRIBUTES)
    for namespace in MANAGED_NAMESPACES:
        await register_namespace_if_missing(env, namespace)
        await ensure_namespace_search_attributes(env, namespace)
    yield env
    await env.shutdown()


@pytest.fixture(scope='session')
def temporal_host(temporal_env):
    return temporal_host_from_env(temporal_env)


@pytest_asyncio.fixture(autouse=True)
async def cleanup_temporal_schedules(temporal_env):
    """Delete orphan schedules in scouter/laborious before each test."""
    host = temporal_host_from_env(temporal_env)
    for namespace in MANAGED_NAMESPACES:
        client = await Client.connect(host, namespace=namespace)
        async for schedule in await client.list_schedules():
            handle = client.get_schedule_handle(schedule.id)
            await handle.delete()
    yield


@pytest_asyncio.fixture
def mock_logger():
    """Logger double with readable console output for E2E runs."""
    logger = MagicMock(spec=Logger)
    logger.info = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
    logger.debug = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
    logger.error = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
    logger.warning = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
    logger.custom_info = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
    logger.custom_debug = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
    logger.custom_error = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
    logger.custom_warning = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
    return logger


@pytest_asyncio.fixture
async def notification_handler(mock_logger, mongo_container):
    """Real notification handler using MongoDB testcontainer."""
    handler = CoreNotificationHandler(
        connection_string=mongo_uri_from_container(mongo_container),
        database=E2E_DATABASE,
        logger=mock_logger,
        project_name='orchestrator-e2e',
    )
    try:
        yield handler
    finally:
        handler.shutdown()


@pytest.fixture
def notification_inserts(notification_handler):
    """Spy on real Mongo insert calls issued by notification handler."""
    collection = notification_handler.mongo_collection
    original_insert_one = collection.insert_one
    spy = MagicMock(wraps=original_insert_one)
    collection.insert_one = spy
    try:
        yield spy
    finally:
        collection.insert_one = original_insert_one


@pytest_asyncio.fixture
async def test_activities(
    postgres_container,
    mongo_container,
    redis_container,
    smtp_server,
    mock_logger,
    notification_handler,
    temporal_host,
):
    """Real Activities wired to testcontainers and in-process SMTP."""
    mongo_port = mongo_container.get_exposed_port(27017)
    redis_port = int(redis_container.get_exposed_port(6379))

    activities = Activities(
        temporal_config={
            'temporal_host': temporal_host,
            'temporal_scouter_namespace': 'scouter',
            'temporal_laborious_namespace': 'laborious',
        },
        redis_config={
            'host': 'localhost',
            'port': redis_port,
            'username': '',
            'password': '',
        },
        mongodb_config={
            'connection_string': f'mongodb://localhost:{mongo_port}',
            'database_name': E2E_DATABASE,
            'ttl_index_seconds': 3600,
        },
        email_config={
            'sender_email': 'e2e@example.com',
            'sender_password': '',
            'smtp_server': smtp_server.host,
            'smtp_port': smtp_server.port,
        },
        postgres_config={
            'host': 'localhost',
            'port': int(postgres_container.get_exposed_port(5432)),
            'user': postgres_container.username,
            'password': postgres_container.password,
            'dbname': postgres_container.dbname,
            'min_connections': 1,
            'max_connections': 5,
        },
        logger=mock_logger,
        notification_handler=notification_handler,
    )
    await activities.connect_to_temporal()
    try:
        yield activities
    finally:
        activities.shutdown()


def _orchestrator_activity_list(activities: Activities) -> list:
    return [
        activities.load_active_ingestors,
        activities.load_opc_slots,
        activities.update_slots,
        activities.delete_slots,
        activities.aggregate_documents_in_mongodb,
        activities.find_documents_in_mongodb,
        activities.update_pipelines_timestamps,
        activities.create_pipelines_timestamps,
        activities.delete_pipelines_timestamps,
        activities.create_collection_with_ttl_index,
        activities.create_schedules,
        activities.update_schedules,
        activities.delete_schedules,
        activities.normalize_schedules,
        activities.process_schedules,
        activities.process_slots,
        activities.create_schedule_config,
        activities.create_slot_config,
        activities.report_schedule_orchestration,
        activities.report_slot_orchestration,
        activities.format_schedule_config,
        activities.get_last_data_timestamp,
        activities.load_latest_data,
        activities.put_last_data_timestamp,
        activities.filter_notification_alerts,
        activities.filter_notification_reports,
        activities.build_email_html,
        activities.send_email,
        activities.format_log_report,
        activities.export_data_to_postgres,
        activities.store_notification_cache,
    ]


@pytest_asyncio.fixture
async def orchestrator_worker(temporal_env, test_activities):
    """Worker for orchestrator workflows and all activities on the default namespace."""
    with ThreadPoolExecutor(max_workers=32) as activity_executor:
        async with Worker(
            temporal_env.client,
            task_queue=ORCHESTRATOR_TASK_QUEUE,
            workflows=[
                Orchestrator,
                Alerts,
                Reports,
                LoadNotificationPackage,
                ProcessNotifications,
            ],
            activities=_orchestrator_activity_list(test_activities),
            activity_executor=activity_executor,
        ) as worker:
            yield worker


@pytest_asyncio.fixture
async def stub_workers(temporal_env):
    """No-op workers on scouter/laborious namespaces for every managed workflow type."""
    host = temporal_host_from_env(temporal_env)
    worker_contexts: list[Worker] = []
    clients: list[Client] = []
    stub_types = list(schedule_types.keys()) + [
        'xgboost_predictions_batch',
        'xgboost_minimal_retrain',
    ]

    try:
        for namespace in MANAGED_NAMESPACES:
            client = await Client.connect(host, namespace=namespace)
            clients.append(client)
            queues = {
                build_queue_name(workflow_type, runtime)
                for workflow_type in stub_types
                for runtime in E2E_RUNTIMES
            }
            for queue in queues:
                worker = Worker(
                    client,
                    task_queue=queue,
                    workflows=STUB_WORKFLOW_CLASSES,
                )
                await worker.__aenter__()
                worker_contexts.append(worker)
        yield worker_contexts
    finally:
        for worker in reversed(worker_contexts):
            await worker.__aexit__(None, None, None)