"""Pipeline testing framework for ETL pipelines.

Provides mock components, assertion helpers, fixture loading, and a test harness
for testing pipelines without real I/O.

Usage:
    from pycharter.etl_generator.testing import (
        MockExtractor, MockLoader, PipelineTestHarness,
        assert_records_match, assert_record_count, load_fixture,
    )

    harness = PipelineTestHarness(pipeline, fixture_data=[{"id": 1}])
    result = await harness.run()
    assert_record_count(result, expected=1)
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, AsyncIterator

import yaml

from pycharter.etl_generator.result import LoadResult, PipelineResult

logger = logging.getLogger(__name__)


# =============================================================================
# MOCK COMPONENTS
# =============================================================================


class MockExtractor:
    """Mock extractor that yields pre-configured fixture data.

    Implements the Extractor protocol for testing pipelines without
    real data sources.

    Args:
        data: Fixture records. Either a flat list of dicts (auto-batched by
            batch_size) or a pre-batched list of lists.
        batch_size: Number of records per batch when data is flat.
    """

    def __init__(
        self,
        data: list[dict[str, Any]] | list[list[dict[str, Any]]],
        batch_size: int = 1000,
    ) -> None:
        self._data = data
        self._batch_size = batch_size

    async def extract(self, **params: Any) -> AsyncIterator[list[dict[str, Any]]]:
        """Yield fixture data as batches.

        Yields:
            Batches of records from the configured fixture data.
        """
        if not self._data:
            return

        # Detect pre-batched vs flat data
        if self._data and isinstance(self._data[0], list):
            for batch in self._data:
                yield batch
        else:
            # Chunk flat data by batch_size
            flat: list[dict[str, Any]] = self._data  # type: ignore[assignment]
            for i in range(0, len(flat), self._batch_size):
                yield flat[i : i + self._batch_size]


@dataclass
class MockLoader:
    """Mock loader that captures loaded records for assertion.

    Implements the Loader protocol. Records are accumulated in
    ``loaded_records`` and each call's batch is stored in ``load_calls``.

    Args:
        simulate_failure: If True, ``load()`` returns a failure result.
        failure_error: Error message used when simulating failure.
    """

    simulate_failure: bool = False
    failure_error: str = "Simulated load failure"
    loaded_records: list[dict[str, Any]] = field(default_factory=list)
    load_calls: list[list[dict[str, Any]]] = field(default_factory=list)

    async def load(self, data: list[dict[str, Any]], **params: Any) -> LoadResult:
        """Capture loaded data and return a result.

        Args:
            data: Batch of records to load.
            **params: Ignored.

        Returns:
            LoadResult indicating success or simulated failure.
        """
        self.load_calls.append(list(data))

        if self.simulate_failure:
            return LoadResult(
                success=False,
                rows_loaded=0,
                rows_failed=len(data),
                error=self.failure_error,
            )

        self.loaded_records.extend(data)
        return LoadResult(success=True, rows_loaded=len(data))

    def reset(self) -> None:
        """Clear all captured records and call history."""
        self.loaded_records.clear()
        self.load_calls.clear()


# =============================================================================
# TEST HARNESS
# =============================================================================


class PipelineTestHarness:
    """Run a pipeline with mock I/O injected.

    Works with both programmatic and config-driven pipelines by replacing
    the extractor and loader with mock implementations.

    Args:
        pipeline: The pipeline to test.
        fixture_data: Test data — flat list of dicts or pre-batched list of lists.
        batch_size: Batch size for the mock extractor.

    Example:
        >>> harness = PipelineTestHarness(
        ...     pipeline, fixture_data=[{"id": 1, "name": "Alice"}]
        ... )
        >>> result = await harness.run()
        >>> assert result.success
        >>> assert harness.loaded_records == [{"id": 1, "name": "Alice"}]
    """

    def __init__(
        self,
        pipeline: Any,
        fixture_data: list[dict[str, Any]] | list[list[dict[str, Any]]],
        batch_size: int = 1000,
    ) -> None:
        self._pipeline = pipeline
        self._mock_extractor = MockExtractor(fixture_data, batch_size=batch_size)
        self._mock_loader = MockLoader()

    @property
    def loaded_records(self) -> list[dict[str, Any]]:
        """Records captured by the mock loader."""
        return self._mock_loader.loaded_records

    @property
    def load_calls(self) -> list[list[dict[str, Any]]]:
        """Individual load call batches."""
        return self._mock_loader.load_calls

    async def run(self, **params: Any) -> PipelineResult:
        """Run the pipeline with mock extractor and loader.

        Args:
            **params: Passed through to ``pipeline.run()``.

        Returns:
            PipelineResult from the pipeline execution.
        """
        self._pipeline.extractor = self._mock_extractor
        self._pipeline.loader = self._mock_loader
        return await self._pipeline.run(**params)


# =============================================================================
# ASSERTION HELPERS
# =============================================================================


def assert_records_match(
    actual: list[dict[str, Any]],
    expected: list[dict[str, Any]],
    *,
    order_matters: bool = True,
    subset: bool = False,
) -> None:
    """Assert that actual records match expected records.

    Args:
        actual: Records produced by the pipeline.
        expected: Expected records.
        order_matters: If True, records must be in the same order.
        subset: If True, actual must contain expected as a subset.

    Raises:
        AssertionError: If records do not match.
    """
    if subset:
        for i, exp in enumerate(expected):
            if exp not in actual:
                raise AssertionError(
                    f"Expected record at index {i} not found in actual: {exp}"
                )
        return

    if not order_matters:
        _assert_unordered(actual, expected)
        return

    if len(actual) != len(expected):
        raise AssertionError(
            f"Record count mismatch: got {len(actual)}, expected {len(expected)}"
        )
    for i, (act, exp) in enumerate(zip(actual, expected)):
        if act != exp:
            raise AssertionError(
                f"Record mismatch at index {i}:\n  actual:   {act}\n  expected: {exp}"
            )


def _assert_unordered(
    actual: list[dict[str, Any]], expected: list[dict[str, Any]]
) -> None:
    """Assert records match regardless of order."""
    if len(actual) != len(expected):
        raise AssertionError(
            f"Record count mismatch: got {len(actual)}, expected {len(expected)}"
        )
    # Convert dicts to sorted-tuple representation for comparison
    actual_set = _to_comparable(actual)
    expected_set = _to_comparable(expected)
    if actual_set != expected_set:
        raise AssertionError(
            f"Records do not match (order-insensitive):\n"
            f"  actual:   {actual}\n  expected: {expected}"
        )


def _to_comparable(records: list[dict[str, Any]]) -> list[tuple[tuple[str, Any], ...]]:
    """Convert records to a sorted list of sorted tuples for comparison."""
    return sorted(tuple(sorted(r.items())) for r in records)


def assert_record_count(result: PipelineResult, expected: int) -> None:
    """Assert the number of loaded rows in a pipeline result.

    Args:
        result: PipelineResult from a pipeline run.
        expected: Expected number of loaded rows.

    Raises:
        AssertionError: If the count does not match.
    """
    if result.rows_loaded != expected:
        raise AssertionError(
            f"Row count mismatch: got {result.rows_loaded}, expected {expected}"
        )


def assert_fields_present(records: list[dict[str, Any]], fields: list[str]) -> None:
    """Assert that all specified fields exist in every record.

    Args:
        records: List of records to check.
        fields: Field names that must be present.

    Raises:
        AssertionError: If any field is missing from any record.
    """
    for i, record in enumerate(records):
        for f in fields:
            if f not in record:
                raise AssertionError(
                    f"Field '{f}' missing from record at index {i}: {record}"
                )


def assert_no_field(records: list[dict[str, Any]], field_name: str) -> None:
    """Assert that a field does not exist in any record.

    Args:
        records: List of records to check.
        field_name: Field name that must be absent.

    Raises:
        AssertionError: If the field is found in any record.
    """
    for i, record in enumerate(records):
        if field_name in record:
            raise AssertionError(
                f"Field '{field_name}' found in record at index {i}: {record}"
            )


def assert_field_values(
    records: list[dict[str, Any]],
    field_name: str,
    expected_values: list[Any],
) -> None:
    """Assert that a specific field has the expected values across records.

    Args:
        records: List of records to check.
        field_name: Field to extract values from.
        expected_values: Expected values in order.

    Raises:
        AssertionError: If values do not match.
    """
    actual_values = [r.get(field_name) for r in records]
    if actual_values != expected_values:
        raise AssertionError(
            f"Field '{field_name}' values mismatch:\n"
            f"  actual:   {actual_values}\n  expected: {expected_values}"
        )


def assert_schema_shape(records: list[dict[str, Any]], schema: dict[str, type]) -> None:
    """Assert that field values match expected types.

    None values are allowed for any field type.

    Args:
        records: List of records to check.
        schema: Mapping of field name to expected Python type.

    Raises:
        AssertionError: If any field has an unexpected type.
    """
    for i, record in enumerate(records):
        for field_name, expected_type in schema.items():
            if field_name not in record:
                raise AssertionError(
                    f"Field '{field_name}' missing from record at index {i}"
                )
            value = record[field_name]
            if value is not None and not isinstance(value, expected_type):
                raise AssertionError(
                    f"Field '{field_name}' at index {i}: expected "
                    f"{expected_type.__name__}, got {type(value).__name__} "
                    f"(value={value!r})"
                )


# =============================================================================
# FIXTURE LOADING
# =============================================================================


@dataclass(frozen=True, slots=True)
class TestFixture:
    """Container for loaded test fixture data.

    Attributes:
        name: Optional fixture name from metadata.
        records: Flat list of records.
        batches: Pre-batched records (empty if data was flat).
        metadata: Additional metadata from the fixture file.
    """

    name: str = ""
    records: tuple[dict[str, Any], ...] = ()
    batches: tuple[tuple[dict[str, Any], ...], ...] = ()
    metadata: dict[str, Any] = field(default_factory=dict)


def load_fixture(path: str | Path) -> list[dict[str, Any]]:
    """Load fixture data from a YAML or JSON file.

    Supported formats:
        - Top-level list: ``[{id: 1}, ...]``
        - Dict with ``records`` key: ``{records: [{id: 1}, ...]}``
        - Dict with ``batches`` key: flattened into a single list.

    Args:
        path: Path to the fixture file (.yaml, .yml, or .json).

    Returns:
        Flat list of records.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the format is unrecognized.
    """
    fixture = load_test_fixture(path)
    return list(fixture.records)


def load_test_fixture(path: str | Path) -> TestFixture:
    """Load a test fixture with full metadata.

    Args:
        path: Path to the fixture file (.yaml, .yml, or .json).

    Returns:
        TestFixture with records, batches, and metadata.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the format is unrecognized.
    """
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Fixture file not found: {path}")

    data = _read_fixture_file(path)
    return _parse_fixture_data(data)


def _read_fixture_file(path: Path) -> Any:
    """Read and parse a fixture file."""
    suffix = path.suffix.lower()
    with open(path) as f:
        if suffix == ".json":
            return json.load(f)
        return yaml.safe_load(f)


def _parse_fixture_data(data: Any) -> TestFixture:
    """Parse raw fixture data into a TestFixture."""
    if isinstance(data, list):
        return TestFixture(records=tuple(data))

    if isinstance(data, dict):
        name = data.get("name", "")
        metadata = {
            k: v for k, v in data.items() if k not in ("name", "records", "batches")
        }

        if "batches" in data:
            batches = tuple(tuple(b) for b in data["batches"])
            flat = [record for batch in data["batches"] for record in batch]
            return TestFixture(
                name=name,
                records=tuple(flat),
                batches=batches,
                metadata=metadata,
            )

        if "records" in data:
            return TestFixture(
                name=name,
                records=tuple(data["records"]),
                metadata=metadata,
            )

    raise ValueError(
        f"Unrecognized fixture format. Expected a list, or a dict with "
        f"'records' or 'batches' key. Got: {type(data).__name__}"
    )


# =============================================================================
# CONFIG VALIDATION HELPER
# =============================================================================


def validate_pipeline_config(
    config: dict[str, Any] | str | Path,
    *,
    variables: dict[str, str] | None = None,
) -> tuple[bool, list[dict[str, Any]]]:
    """Validate a pipeline configuration for correctness.

    Wraps the existing ``ConfigValidator`` for a simpler testing API.

    Args:
        config: Pipeline config as a dict, file path string, or Path object.
        variables: Optional variables for ``${VAR}`` substitution.

    Returns:
        Tuple of (is_valid, errors) where errors is a list of dicts
        with 'section' and 'message' keys.
    """
    from pycharter.etl_generator.config_validator import ConfigValidator
    from pycharter.etl_generator.context import PipelineContext

    errors: list[dict[str, Any]] = []

    # Load config from file if needed
    if isinstance(config, (str, Path)):
        path = Path(config)
        if not path.exists():
            return False, [{"section": "file", "message": f"File not found: {path}"}]
        if path.is_dir():
            return _validate_config_dir(path, variables)
        with open(path) as f:
            raw = f.read()
        if variables:
            context = PipelineContext(variables=variables)
            raw = context.resolve(raw)
        config = yaml.safe_load(raw) or {}

    if not isinstance(config, dict):
        return False, [{"section": "config", "message": "Config must be a dict"}]

    validator = ConfigValidator(strict=True)

    # Validate extract
    extract_config = config.get("extract", {})
    if extract_config:
        is_valid, extract_errors = validator.check_extract(extract_config)
        if not is_valid:
            for err in extract_errors:
                errors.append({"section": "extract", "message": str(err)})

    # Validate transform
    transform_config = config.get("transform")
    if transform_config:
        if isinstance(transform_config, list):
            is_valid, transform_errors = validator.check_transform(
                {"transform": transform_config}
            )
        else:
            is_valid, transform_errors = validator.check_transform(transform_config)
        if not is_valid:
            for err in transform_errors:
                errors.append({"section": "transform", "message": str(err)})

    # Validate load
    load_config = config.get("load", {})
    if load_config:
        is_valid, load_errors = validator.check_load(load_config)
        if not is_valid:
            for err in load_errors:
                errors.append({"section": "load", "message": str(err)})

    return len(errors) == 0, errors


def _validate_config_dir(
    path: Path, variables: dict[str, str] | None
) -> tuple[bool, list[dict[str, Any]]]:
    """Validate a config directory with extract.yaml, load.yaml, etc."""
    from pycharter.etl_generator.config_validator import ConfigValidator
    from pycharter.etl_generator.context import PipelineContext

    errors: list[dict[str, Any]] = []
    validator = ConfigValidator(strict=True)
    context = PipelineContext(variables=variables) if variables else None

    for name, required in [("extract", True), ("load", True), ("transform", False)]:
        file = path / f"{name}.yaml"
        if not file.exists():
            if required:
                errors.append(
                    {"section": name, "message": f"Required file missing: {file}"}
                )
            continue

        with open(file) as f:
            raw = f.read()
        if context:
            raw = context.resolve(raw)
        cfg = yaml.safe_load(raw) or {}

        check_method = getattr(validator, f"check_{name}")
        if name == "transform" and isinstance(cfg, list):
            is_valid, check_errors = check_method({"transform": cfg})
        else:
            is_valid, check_errors = check_method(cfg)
        if not is_valid:
            for err in check_errors:
                errors.append({"section": name, "message": str(err)})

    return len(errors) == 0, errors
