"""Test utilities for HLA-Compass module developers.

Provides helpers for writing unit tests without needing a real API connection
or cloud infrastructure.

Usage::

    from hla_compass.testing import make_test_context, MockStorageClient, run_module

    def test_my_module():
        result = run_module(MyModule, {"sequence": "SIINFEKL"})
        assert result["status"] == "success"

    def test_with_storage():
        storage = MockStorageClient()
        result = run_module(MyModule, {"x": 1}, storage=storage)
        assert "output.json" in storage.objects
"""

from __future__ import annotations

import json
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Type


def make_test_context(**overrides: Any) -> Dict[str, Any]:
    """Create a minimal valid execution context dict for testing.

    All fields have sensible defaults.  Pass keyword arguments to override
    specific values.  The context has ``offline=True`` by default so no
    real API calls are attempted.

    Returns:
        A dict compatible with :class:`~hla_compass.context.RuntimeContext`.
    """
    base: Dict[str, Any] = {
        "run_id": f"test-{uuid.uuid4().hex[:8]}",
        "module_id": "test-module",
        "module_version": "1.0.0",
        "organization_id": "test-org",
        "user_id": "test-user",
        "environment": "test",
        "correlation_id": f"test-{uuid.uuid4().hex[:8]}",
        "requested_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "roles": ["developer"],
        "mode": "interactive",
        "offline": True,
    }
    base.update(overrides)
    return base


class MockStorageClient:
    """In-memory storage backend that captures uploads.

    Useful for verifying that a module writes the expected files
    without touching S3.

    Attributes:
        objects: Mapping of key → raw bytes for each saved object.
        metadata: Mapping of key → metadata dict.
    """

    def __init__(self, *, bucket: str = "mock-bucket", region: str = "us-east-1"):
        self.bucket = bucket
        self.region = region
        self.objects: Dict[str, bytes] = {}
        self.metadata: Dict[str, Dict[str, str]] = {}

    # --- write interface (matches S3StorageClient) ---

    def put_object(
        self,
        key: str,
        body: bytes,
        content_type: str = "application/octet-stream",
        metadata: Optional[Dict[str, str]] = None,
    ) -> str:
        self.objects[key] = body
        self.metadata[key] = {"content_type": content_type, **(metadata or {})}
        return f"mock://{self.bucket}/{key}"

    # --- read interface ---

    def get_object(self, key: str) -> bytes:
        if key not in self.objects:
            raise KeyError(f"Object not found: {key}")
        return self.objects[key]

    def list_objects(self, prefix: Optional[str] = None) -> List[Dict[str, Any]]:
        return [{"key": k, "size": len(v)} for k, v in self.objects.items() if prefix is None or k.startswith(prefix)]

    def delete_object(self, key: str) -> None:
        self.objects.pop(key, None)
        self.metadata.pop(key, None)

    def create_presigned_url(self, key: str, expires_in: int = 3600) -> str:
        return f"mock://presigned/{self.bucket}/{key}?expires={expires_in}"

    # --- convenience ---

    def get_json(self, key: str) -> Any:
        """Parse a stored object as JSON."""
        return json.loads(self.get_object(key))


class MockSQLClient:
    """Mock SQL client that returns canned query results.

    Args:
        results_map: Mapping of SQL substring → result dict.
            When a query contains a key from this map, the corresponding
            result is returned.  Falls back to an empty result set.
    """

    def __init__(self, results_map: Optional[Dict[str, Dict[str, Any]]] = None):
        self._results = results_map or {}
        self.queries: List[tuple] = []

    def query(self, sql: str, params: Optional[List[Any]] = None) -> Dict[str, Any]:
        self.queries.append((sql, params))
        for pattern, result in self._results.items():
            if pattern in sql:
                return result
        return {"columns": [], "data": [], "count": 0}

    def query_df(self, sql: str, params: Optional[List[Any]] = None, engine: str = "pandas"):
        result = self.query(sql, params)
        columns = result.get("columns", [])
        data = result.get("data", [])

        if engine == "pandas":
            try:
                import pandas as pd
            except ImportError as err:
                raise ImportError("pandas is required for query_df()") from err
            return pd.DataFrame(data, columns=columns if columns else None)

        if engine == "polars":
            try:
                import polars as pl
            except ImportError as err:
                raise ImportError("polars is required for query_df()") from err
            if columns:
                return pl.DataFrame(data, schema=columns, orient="row")
            return pl.DataFrame(data)

        raise ValueError(f"Unknown engine: {engine!r}")


class MockDataClient:
    """Mock DataClient with a canned SQL client.

    Args:
        query_results: Passed to :class:`MockSQLClient`.
    """

    def __init__(self, query_results: Optional[Dict[str, Dict[str, Any]]] = None):
        self.sql = MockSQLClient(query_results)
        self.storage = None
        self.catalog_storage = None


def run_module(
    module_class: Type,
    input_data: Dict[str, Any],
    *,
    manifest_path: Optional[str] = None,
    storage: Optional[MockStorageClient] = None,
    **context_overrides: Any,
) -> Dict[str, Any]:
    """Instantiate and execute a module in one call.

    Args:
        module_class: The Module subclass to test.
        input_data: Input parameters for the module.
        manifest_path: Path to manifest.json (defaults to CWD/manifest.json).
        storage: Optional :class:`MockStorageClient` to inject.
        **context_overrides: Extra fields merged into the test context.

    Returns:
        The module output dict (with ``status``, ``results``, ``metadata``, etc.).
    """
    module = module_class(manifest_path=manifest_path)
    ctx = make_test_context(**context_overrides)
    if storage is not None:
        ctx["storage"] = storage
    return module.run(input_data, ctx)
