from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any, Dict, Optional

import pytest

# Attempt to import project modules. The tests are resilient:
# - If the real slash engine API exists, we'll try to use it.
# - Otherwise we fall back to a local reference implementation so the tests still run.

try:
    from aicf.sla import slash_engine as se  # type: ignore
except Exception:  # pragma: no cover - optional path
    se = None  # type: ignore

try:
    from aicf.sla import evaluator as ev  # type: ignore
except Exception:  # pragma: no cover - optional path
    ev = None  # type: ignore

try:
    from aicf.aitypes import sla as sla_types  # type: ignore
except Exception:  # pragma: no cover - optional path
    sla_types = None  # type: ignore


# --------------------------- Local, deterministic fallbacks ---------------------------


@dataclass
class _Thresholds:
    traps_min: float = 0.98
    qos_min: float = 0.90


@dataclass
class _PenaltyRules:
    penalty_per_violation: int = 1_000
    jail_after_violations: int = 2
    cooldown_blocks: int = 5


@dataclass
class _Provider:
    provider_id: str
    stake: int
    jailed: bool = False
    jail_until_height: int = 0
    violations: int = 0


def _z_from_conf(conf: float) -> float:
    table = {0.80: 1.2816, 0.90: 1.6449, 0.95: 1.96, 0.975: 2.2414, 0.99: 2.5758}
    closest = min(table, key=lambda k: abs(k - conf))
    return table[closest]


def _wilson_lower(successes: int, total: int, conf: float) -> float:
    if total <= 0:
        return 0.0
    z = _z_from_conf(conf)
    phat = successes / total
    denom = 1 + (z * z) / total
    center = phat + (z * z) / (2 * total)
    margin = z * math.sqrt((phat * (1 - phat) + (z * z) / (4 * total)) / total)
    return (center - margin) / denom


def _eval_dimensions(
    stats: Dict[str, int], thresholds: _Thresholds, conf: float
) -> Dict[str, bool]:
    total = int(stats.get("total", 0))
    traps_ok = int(stats.get("traps_ok", 0))
    qos_ok = int(stats.get("qos_ok", 0))
    traps_pass = _wilson_lower(traps_ok, total, conf) >= thresholds.traps_min
    qos_pass = _wilson_lower(qos_ok, total, conf) >= thresholds.qos_min
    return {"traps": traps_pass, "qos": qos_pass, "overall": traps_pass and qos_pass}


class _LocalSlashEngine:
    def __init__(
        self, penalties: _PenaltyRules, thresholds: _Thresholds, conf: float = 0.95
    ):
        self.penalties = penalties
        self.thresholds = thresholds
        self.conf = conf

    def process_window(
        self, provider: _Provider, height: int, stats: Dict[str, int]
    ) -> Optional[Dict[str, Any]]:
        # If jailed and cooldown not elapsed, do nothing.
        if provider.jailed and height < provider.jail_until_height:
            return None

        res = _eval_dimensions(stats, self.thresholds, self.conf)
        if res["overall"]:
            # Good window. If past cooldown, auto-unjail.
            if provider.jailed and height >= provider.jail_until_height:
                provider.jailed = False
                provider.violations = 0
            return None

        # Bad window: apply slash.
        provider.violations += 1
        penalty = self.penalties.penalty_per_violation
        provider.stake = max(0, provider.stake - penalty)

        event = {
            "kind": "SlashEvent",
            "provider_id": provider.provider_id,
            "penalty": penalty,
            "height": height,
            "violations": provider.violations,
            "jailed": False,
        }

        if provider.violations >= self.penalties.jail_after_violations:
            provider.jailed = True
            provider.jail_until_height = height + self.penalties.cooldown_blocks
            event["jailed"] = True
            event["jail_until_height"] = provider.jail_until_height

        return event


# --------------------------- Adapters into project modules (best-effort) ---------------------------


def _get_thresholds() -> _Thresholds:
    if sla_types is None:
        return _Thresholds()
    for name in ("Thresholds", "SlaThresholds", "PolicyThresholds"):
        if hasattr(sla_types, name):
            T = getattr(sla_types, name)
            try:
                t = T()  # type: ignore[call-arg]
                traps = float(getattr(t, "traps_min", 0.98))
                qos = float(getattr(t, "qos_min", 0.90))
                return _Thresholds(traps_min=traps, qos_min=qos)
            except Exception:
                continue
    return _Thresholds()


def _maybe_project_engine() -> Optional[Any]:
    """Try to construct the project's slash engine. Return None if not possible."""
    if se is None:
        return None
    # Common shapes:
    #   - se.SlashEngine(thresholds=..., penalties=..., conf=...)
    #   - se.SlashEngine(...)
    #   - module-level process/evaluate functions
    for name in ("SlashEngine", "Engine", "Evaluator"):
        if hasattr(se, name):
            C = getattr(se, name)
            try:
                return C()  # type: ignore[call-arg]
            except Exception:
                try:
                    return C(thresholds=_get_thresholds())  # type: ignore[call-arg]
                except Exception:
                    try:
                        return C(penalties=_PenaltyRules(), thresholds=_get_thresholds())  # type: ignore[call-arg]
                    except Exception:
                        pass
    # Fallback to module-level function wrappers
    for fn in ("process_window", "evaluate_and_maybe_slash", "maybe_slash"):
        if hasattr(se, fn):
            return getattr(se, fn)
    return None


def _process_with_engine(
    engine: Any, provider: _Provider, height: int, stats: Dict[str, int]
) -> Optional[Dict[str, Any]]:
    """Normalize processing across possible project engine APIs. If not supported, return None to let caller fallback."""
    if engine is None:
        return None

    # If engine is a function
    if callable(engine) and not hasattr(engine, "process_window"):
        try:
            res = engine(provider=provider, height=height, stats=stats)  # type: ignore[misc]
            return _normalize_slash_event(res)
        except TypeError:
            try:
                res = engine(provider, stats, height)  # type: ignore[misc]
                return _normalize_slash_event(res)
            except Exception:
                return None
        except Exception:
            return None

    # If engine is an object
    for method in (
        "process_window",
        "evaluate_and_maybe_slash",
        "maybe_slash",
        "on_window",
    ):
        if hasattr(engine, method):
            try:
                res = getattr(engine, method)(provider=provider, height=height, stats=stats)  # type: ignore[misc]
                return _normalize_slash_event(res)
            except TypeError:
                try:
                    res = getattr(engine, method)(provider, stats, height)  # type: ignore[misc]
                    return _normalize_slash_event(res)
                except Exception:
                    continue
            except Exception:
                continue
    return None


def _normalize_slash_event(res: Any) -> Optional[Dict[str, Any]]:
    if res is None:
        return None
    if isinstance(res, dict):
        return res
    # Try object with attributes
    out: Dict[str, Any] = {}
    for k in (
        "kind",
        "provider_id",
        "penalty",
        "height",
        "violations",
        "jailed",
        "jail_until_height",
    ):
        if hasattr(res, k):
            out[k] = getattr(res, k)
    return out or None


# --------------------------- Fixtures ---------------------------


@pytest.fixture
def thresholds() -> _Thresholds:
    return _get_thresholds()


@pytest.fixture
def penalties() -> _PenaltyRules:
    return _PenaltyRules()


@pytest.fixture
def local_engine(
    thresholds: _Thresholds, penalties: _PenaltyRules
) -> _LocalSlashEngine:
    return _LocalSlashEngine(penalties=penalties, thresholds=thresholds, conf=0.95)


@pytest.fixture
def maybe_proj_engine() -> Optional[Any]:
    return _maybe_project_engine()


@pytest.fixture
def provider() -> _Provider:
    return _Provider(provider_id="prov-1", stake=10_000)


# --------------------------- Tests ---------------------------


def test_penalty_and_jail_on_repeated_failures(
    provider: _Provider,
    thresholds: _Thresholds,
    penalties: _PenaltyRules,
    local_engine: _LocalSlashEngine,
    maybe_proj_engine: Optional[Any],
) -> None:
    """
    Two consecutive failing windows should:
      - deduct stake twice (>= 2 * penalty_per_violation total)
      - jail the provider on/after the second violation
    """
    height = 1
    bad_stats = {
        "total": 200,
        "traps_ok": 190,
        "qos_ok": 150,
    }  # clearly below thresholds

    # First failure
    ev1 = _process_with_engine(maybe_proj_engine, provider, height, bad_stats)
    if ev1 is None:
        ev1 = local_engine.process_window(provider, height, bad_stats)

    assert ev1 is not None, "Expected a slash event on first bad window"
    assert provider.stake <= 9_000, "Stake must be reduced after first violation"
    jailed_after_first = provider.jailed

    # Second failure
    height += 1
    ev2 = _process_with_engine(maybe_proj_engine, provider, height, bad_stats)
    if ev2 is None:
        ev2 = local_engine.process_window(provider, height, bad_stats)

    assert ev2 is not None, "Expected a slash event on second bad window"
    assert provider.stake <= 8_000, "Stake must be reduced again after second violation"
    assert provider.violations >= 2, "Violation counter should reflect repeats"
    assert provider.jailed, "Provider should be jailed after repeated violations"
    assert provider.jail_until_height >= height, "Jail should set a cooldown end height"
    # Ensure jailing transitioned (if not jailed after first, it must be jailed now)
    assert (
        not jailed_after_first
    ) or provider.jailed, "Provider should be jailed by now"


def test_recovery_after_cooldown_and_topup(
    provider: _Provider,
    thresholds: _Thresholds,
    penalties: _PenaltyRules,
    local_engine: _LocalSlashEngine,
    maybe_proj_engine: Optional[Any],
) -> None:
    """
    After being jailed, a provider should recover when:
      - cooldown elapses, and
      - it delivers a good window (or the engine unjails automatically post-cooldown).
    Stake top-ups are allowed any time (simulated here).
    """
    # Force jail with two failures
    height = 10
    bad_stats = {"total": 200, "traps_ok": 180, "qos_ok": 150}
    _ = _process_with_engine(
        maybe_proj_engine, provider, height, bad_stats
    ) or local_engine.process_window(provider, height, bad_stats)
    height += 1
    _ = _process_with_engine(
        maybe_proj_engine, provider, height, bad_stats
    ) or local_engine.process_window(provider, height, bad_stats)

    assert provider.jailed, "Provider should be jailed after consecutive failures"
    cooldown_end = provider.jail_until_height

    # While jailed and before cooldown ends, processing should do nothing (no slash/no unjail).
    height = cooldown_end - 1
    ok_stats = {"total": 200, "traps_ok": 199, "qos_ok": 190}
    ev_mid = _process_with_engine(maybe_proj_engine, provider, height, ok_stats)
    if ev_mid is None:
        ev_mid = local_engine.process_window(provider, height, ok_stats)
    assert ev_mid is None, "No action expected before cooldown end"
    assert provider.jailed, "Still jailed before cooldown end"

    # Top-up stake while jailed (simulated external action)
    provider.stake += 5_000

    # After cooldown, next good window should clear jail (or engine should auto-clear).
    height = cooldown_end
    ev_post = _process_with_engine(maybe_proj_engine, provider, height, ok_stats)
    if ev_post is None:
        ev_post = local_engine.process_window(provider, height, ok_stats)

    # Either auto-unjailed or unjailed due to good window
    assert (
        provider.jailed is False
    ), "Provider should be unjailed after cooldown with good performance"
    assert (
        provider.stake >= 5_000
    ), "Stake top-up must persist (no penalty on good window)"
    assert provider.violations in (
        0,
        1,
        2,
    ), "Engine may reset or keep history; only check bounds"

    # Subsequent good window must not re-jail or slash
    height += 1
    stake_before = provider.stake
    ev_ok2 = _process_with_engine(maybe_proj_engine, provider, height, ok_stats)
    if ev_ok2 is None:
        ev_ok2 = local_engine.process_window(provider, height, ok_stats)
    assert (
        provider.jailed is False
    ), "Must remain unjailed on continued good performance"
    assert provider.stake == stake_before, "No slashing on good windows"
