"""Chaos engineering for flaky test reproduction.

Uses Toxiproxy to inject controlled faults:
- Network latency (slow_close, latency toxics)
- Connection resets (reset_peer toxic)
- Bandwidth limits (bandwidth toxic)
- Timeouts (timeout toxic)

GitHub: https://github.com/Shopify/toxiproxy

Installation:
    brew install toxiproxy  # macOS
    docker run -d -p 8474:8474 shopify/toxiproxy  # Docker
    pip install toxiproxy-python  # Python client
"""
from __future__ import annotations

import logging
import os
import threading
from collections.abc import AsyncGenerator, Generator
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Try to import toxiproxy
_TOXIPROXY_AVAILABLE = False

try:
    from toxiproxy import Toxiproxy
    from toxiproxy.toxic import Toxic

    _TOXIPROXY_AVAILABLE = True
except ImportError:
    Toxiproxy = None
    Toxic = None


class ChaosType(str, Enum):
    """Types of chaos that can be injected."""

    LATENCY = "latency"  # Add network latency
    BANDWIDTH = "bandwidth"  # Limit bandwidth
    SLOW_CLOSE = "slow_close"  # Delay connection close
    TIMEOUT = "timeout"  # Force timeouts
    RESET_PEER = "reset_peer"  # Reset connections
    SLICER = "slicer"  # Slice data into smaller bits
    LIMIT_DATA = "limit_data"  # Limit data transfer


@dataclass
class ChaosConfig:
    """Configuration for chaos injection.

    Attributes:
        chaos_type: Type of fault to inject
        target_service: Service name to target (e.g., "postgres", "redis")
        probability: Probability of fault (0.0-1.0)
        latency_ms: Latency to add (for LATENCY type)
        jitter_ms: Random jitter to add to latency
        bandwidth_kb: Bandwidth limit in KB/s (for BANDWIDTH type)
        timeout_ms: Timeout duration (for TIMEOUT type)
        duration_seconds: How long to inject chaos
    """

    chaos_type: ChaosType
    target_service: str = "database"
    probability: float = 1.0
    latency_ms: int = 100
    jitter_ms: int = 50
    bandwidth_kb: int = 10
    timeout_ms: int = 5000
    duration_seconds: float = 30.0
    attributes: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "chaos_type": self.chaos_type.value,
            "target_service": self.target_service,
            "probability": self.probability,
            "latency_ms": self.latency_ms,
            "jitter_ms": self.jitter_ms,
            "bandwidth_kb": self.bandwidth_kb,
            "timeout_ms": self.timeout_ms,
            "duration_seconds": self.duration_seconds,
            "attributes": self.attributes,
        }

    @classmethod
    def for_timing_flake(cls) -> ChaosConfig:
        """Create config that reproduces timing-related flakes."""
        return cls(
            chaos_type=ChaosType.LATENCY,
            latency_ms=200,
            jitter_ms=100,
            probability=0.8,
        )

    @classmethod
    def for_timeout_flake(cls) -> ChaosConfig:
        """Create config that reproduces timeout-related flakes."""
        return cls(
            chaos_type=ChaosType.TIMEOUT,
            timeout_ms=10000,
            probability=0.5,
        )

    @classmethod
    def for_connection_flake(cls) -> ChaosConfig:
        """Create config that reproduces connection-related flakes."""
        return cls(
            chaos_type=ChaosType.RESET_PEER,
            probability=0.3,
        )

    @classmethod
    def for_slow_network(cls) -> ChaosConfig:
        """Create config that simulates slow network."""
        return cls(
            chaos_type=ChaosType.BANDWIDTH,
            bandwidth_kb=5,
            probability=1.0,
        )


class ToxiproxyClient:
    """Client for Toxiproxy server.

    Toxiproxy is a TCP proxy that allows injecting faults like:
    - Latency
    - Bandwidth limits
    - Connection resets
    - Timeouts

    Docs: https://github.com/Shopify/toxiproxy
    """

    DEFAULT_HOST = "localhost"
    DEFAULT_PORT = 8474

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
    ):
        """Initialize Toxiproxy client.

        Args:
            host: Toxiproxy server host
            port: Toxiproxy server port
        """
        self.host = host or os.environ.get("TOXIPROXY_HOST", self.DEFAULT_HOST)
        self.port = port or int(os.environ.get("TOXIPROXY_PORT", self.DEFAULT_PORT))
        self._client: Optional[Any] = None
        self._proxies: Dict[str, Any] = {}

    @property
    def available(self) -> bool:
        """Check if Toxiproxy is available."""
        if not _TOXIPROXY_AVAILABLE:
            return False
        try:
            self._get_client()
            return True
        except Exception:
            return False

    def _get_client(self) -> Any:
        """Get or create Toxiproxy client."""
        if not _TOXIPROXY_AVAILABLE:
            raise RuntimeError("toxiproxy-python not installed. Run: pip install toxiproxy-python")

        if self._client is None:
            self._client = Toxiproxy(server_host=f"{self.host}:{self.port}")
        return self._client

    def create_proxy(
        self,
        name: str,
        listen_port: int,
        upstream_host: str,
        upstream_port: int,
    ) -> Any:
        """Create a new proxy.

        Args:
            name: Proxy name
            listen_port: Port to listen on
            upstream_host: Upstream service host
            upstream_port: Upstream service port

        Returns:
            Toxiproxy proxy object
        """
        client = self._get_client()

        # Check if proxy already exists
        try:
            existing = client.get_proxy(name)
            if existing:
                return existing
        except Exception:
            pass

        proxy = client.create(
            name=name,
            upstream=f"{upstream_host}:{upstream_port}",
            listen=f"localhost:{listen_port}",
        )
        self._proxies[name] = proxy
        return proxy

    def add_toxic(
        self,
        proxy_name: str,
        toxic_type: str,
        attributes: Dict[str, Any],
        stream: str = "downstream",
        toxicity: float = 1.0,
    ) -> Any:
        """Add a toxic to a proxy.

        Args:
            proxy_name: Name of proxy to add toxic to
            toxic_type: Type of toxic (latency, bandwidth, etc.)
            attributes: Toxic-specific attributes
            stream: "downstream" or "upstream"
            toxicity: Probability of toxic taking effect (0.0-1.0)

        Returns:
            Toxic object
        """
        client = self._get_client()
        proxy = client.get_proxy(proxy_name)

        return proxy.add_toxic(
            type=toxic_type,
            stream=stream,
            toxicity=toxicity,
            attributes=attributes,
        )

    def remove_toxic(self, proxy_name: str, toxic_name: str) -> None:
        """Remove a toxic from a proxy.

        Args:
            proxy_name: Name of proxy
            toxic_name: Name of toxic to remove
        """
        client = self._get_client()
        proxy = client.get_proxy(proxy_name)
        toxic = proxy.get_toxic(toxic_name)
        toxic.destroy()

    def clear_toxics(self, proxy_name: str) -> None:
        """Remove all toxics from a proxy.

        Args:
            proxy_name: Name of proxy
        """
        client = self._get_client()
        proxy = client.get_proxy(proxy_name)

        for toxic in proxy.toxics():
            toxic.destroy()

    def delete_proxy(self, name: str) -> None:
        """Delete a proxy.

        Args:
            name: Proxy name
        """
        client = self._get_client()
        proxy = client.get_proxy(name)
        proxy.destroy()
        self._proxies.pop(name, None)


class ChaosController:
    """Controller for injecting chaos into tests.

    Usage:
        controller = ChaosController()

        # Inject latency to database connections
        with controller.inject(ChaosConfig.for_timing_flake()):
            run_my_test()

        # Or async
        async with controller.inject_async(config):
            await run_my_async_test()
    """

    def __init__(self, toxiproxy_client: Optional[ToxiproxyClient] = None):
        """Initialize chaos controller.

        Args:
            toxiproxy_client: Optional Toxiproxy client (created if not provided)
        """
        self.client = toxiproxy_client or ToxiproxyClient()
        self._active_toxics: List[Dict[str, str]] = []

    @property
    def available(self) -> bool:
        """Check if chaos injection is available."""
        return self.client.available

    @contextmanager
    def inject(self, config: Optional[ChaosConfig] = None) -> Generator[None, None, None]:
        """Inject chaos for the duration of a context.

        If *config* is ``None``, falls back to ``_default_config`` set by
        the ``inject_chaos()`` convenience function (P8-F11).

        When ``config.duration_seconds`` is set to a positive value, the
        context manager will automatically stop chaos after that many
        seconds even if the caller has not exited the ``with`` block
        (P8-F8).

        Args:
            config: Chaos configuration (uses _default_config if None)

        Yields:
            None (chaos is active during context)
        """
        if config is None:
            config = getattr(self, "_default_config", None) or ChaosConfig.for_timing_flake()

        toxic_info = None
        timer_handle: Optional[threading.Timer] = None
        try:
            toxic_info = self._start_chaos(config)

            # P8-F8: Enforce duration_seconds with a background timer
            if toxic_info and config.duration_seconds > 0:
                def _auto_stop():
                    logger.info(
                        f"Chaos duration ({config.duration_seconds}s) expired, "
                        f"auto-removing toxic from {toxic_info['proxy']}"
                    )
                    self._stop_chaos(toxic_info)

                timer_handle = threading.Timer(config.duration_seconds, _auto_stop)
                timer_handle.daemon = True
                timer_handle.start()

            yield
        finally:
            if timer_handle is not None:
                timer_handle.cancel()
            if toxic_info and toxic_info in self._active_toxics:
                self._stop_chaos(toxic_info)

    @asynccontextmanager
    async def inject_async(self, config: ChaosConfig) -> AsyncGenerator[None, None]:
        """Inject chaos for async context.

        Args:
            config: Chaos configuration

        Yields:
            None (chaos is active during context)
        """
        toxic_info = None
        try:
            toxic_info = self._start_chaos(config)
            yield
        finally:
            if toxic_info:
                self._stop_chaos(toxic_info)

    def _start_chaos(self, config: ChaosConfig) -> Optional[Dict[str, str]]:
        """Start chaos injection.

        Args:
            config: Chaos configuration

        Returns:
            Dict with proxy and toxic names for cleanup
        """
        if not self.available:
            logger.warning("Toxiproxy not available, chaos injection skipped")
            return None

        try:
            proxy_name = f"hikaflow-{config.target_service}"
            toxic_name = f"{config.chaos_type.value}-{id(config)}"

            # Map chaos type to toxic type and attributes
            toxic_type, attributes = self._get_toxic_params(config)

            self.client.add_toxic(
                proxy_name=proxy_name,
                toxic_type=toxic_type,
                attributes=attributes,
                toxicity=config.probability,
            )

            info = {"proxy": proxy_name, "toxic": toxic_name}
            self._active_toxics.append(info)

            logger.info(
                f"Chaos injected: {config.chaos_type.value} on {config.target_service}"
            )
            return info

        except Exception as e:
            logger.error(f"Failed to inject chaos: {e}")
            return None

    def _stop_chaos(self, toxic_info: Dict[str, str]) -> None:
        """Stop chaos injection.

        Args:
            toxic_info: Dict with proxy and toxic names
        """
        try:
            self.client.remove_toxic(
                proxy_name=toxic_info["proxy"],
                toxic_name=toxic_info["toxic"],
            )
            self._active_toxics.remove(toxic_info)
            logger.info(f"Chaos removed from {toxic_info['proxy']}")
        except Exception as e:
            logger.warning(f"Failed to remove chaos: {e}")

    def _get_toxic_params(self, config: ChaosConfig) -> tuple[str, Dict[str, Any]]:
        """Get toxic type and attributes from config.

        Args:
            config: Chaos configuration

        Returns:
            Tuple of (toxic_type, attributes)
        """
        if config.chaos_type == ChaosType.LATENCY:
            return "latency", {
                "latency": config.latency_ms,
                "jitter": config.jitter_ms,
            }
        elif config.chaos_type == ChaosType.BANDWIDTH:
            return "bandwidth", {
                "rate": config.bandwidth_kb,
            }
        elif config.chaos_type == ChaosType.SLOW_CLOSE:
            return "slow_close", {
                "delay": config.latency_ms,
            }
        elif config.chaos_type == ChaosType.TIMEOUT:
            return "timeout", {
                "timeout": config.timeout_ms,
            }
        elif config.chaos_type == ChaosType.RESET_PEER:
            return "reset_peer", {
                "timeout": config.timeout_ms,
            }
        elif config.chaos_type == ChaosType.SLICER:
            return "slicer", {
                "average_size": 1024,
                "size_variation": 512,
                "delay": 10,
            }
        elif config.chaos_type == ChaosType.LIMIT_DATA:
            return "limit_data", {
                "bytes": config.bandwidth_kb * 1024,
            }
        else:
            return "latency", {"latency": 100}

    def cleanup_all(self) -> None:
        """Remove all active toxics."""
        for info in list(self._active_toxics):
            self._stop_chaos(info)


def inject_chaos(
    chaos_type: str = "latency",
    target: str = "database",
    probability: float = 1.0,
    **kwargs: Any,
) -> ChaosController:
    """Convenience function to create a chaos controller with config.

    Args:
        chaos_type: Type of chaos (latency, timeout, reset_peer, etc.)
        target: Target service name
        probability: Probability of fault
        **kwargs: Additional config options (latency_ms, jitter_ms, etc.)

    Returns:
        ChaosController instance (use with controller.inject(config))

    Usage:
        controller = inject_chaos("latency", latency_ms=200)
        config = ChaosConfig(
            chaos_type=ChaosType(chaos_type),
            target_service=target,
            probability=probability,
            **{k: v for k, v in kwargs.items() if k in ChaosConfig.__dataclass_fields__},
        )
        with controller.inject(config):
            test_my_function()
    """
    controller = ChaosController()
    # Store the config params for convenience access
    try:
        controller._default_config = ChaosConfig(
            chaos_type=ChaosType(chaos_type),
            target_service=target,
            probability=probability,
            **{k: v for k, v in kwargs.items() if k in ChaosConfig.__dataclass_fields__},
        )
    except (ValueError, TypeError):
        controller._default_config = ChaosConfig.for_timing_flake()
    return controller


# Pytest fixture for chaos testing
def pytest_chaos_fixture():
    """Pytest fixture for chaos testing.

    Usage in conftest.py:
        from autodebug.chaos import pytest_chaos_fixture
        chaos = pytest_chaos_fixture()

    Usage in tests:
        def test_handles_latency(chaos):
            with chaos.inject(ChaosConfig.for_timing_flake()):
                result = my_function()
                assert result is not None
    """
    import pytest

    @pytest.fixture
    def chaos():
        controller = ChaosController()
        yield controller
        controller.cleanup_all()

    return chaos


# Convenience presets for common flaky scenarios
CHAOS_PRESETS = {
    "timing": ChaosConfig.for_timing_flake(),
    "timeout": ChaosConfig.for_timeout_flake(),
    "connection": ChaosConfig.for_connection_flake(),
    "slow_network": ChaosConfig.for_slow_network(),
}


def get_preset(name: str) -> ChaosConfig:
    """Get a preset chaos configuration.

    Args:
        name: Preset name (timing, timeout, connection, slow_network)

    Returns:
        ChaosConfig for the preset
    """
    return CHAOS_PRESETS.get(name, ChaosConfig.for_timing_flake())
