"""Robot Framework helper for ACMS Pipeline Orchestrator integration tests.

Provides a CLI-style interface for Robot to invoke Phase 1 pipeline
component operations and verify the results.

Usage:
    python robot/helper_acms_pipeline_orchestrator.py cws-select
    python robot/helper_acms_pipeline_orchestrator.py cws-preferred
    python robot/helper_acms_pipeline_orchestrator.py pba-proportional
    python robot/helper_acms_pipeline_orchestrator.py pba-min-budget
    python robot/helper_acms_pipeline_orchestrator.py cb-open
    python robot/helper_acms_pipeline_orchestrator.py cb-reset
    python robot/helper_acms_pipeline_orchestrator.py executor-run
    python robot/helper_acms_pipeline_orchestrator.py pipeline-assemble
    python robot/helper_acms_pipeline_orchestrator.py pipeline-custom
"""

from __future__ import annotations

import sys
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Any

_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from cleveragents.application.services.acms_pipeline import (  # noqa: E402
    CircuitBreaker,
    ConfidenceWeightedSelector,
    ContextAssemblyPipeline,
    ParallelStrategyExecutor,
    ProportionalBudgetAllocator,
)
from cleveragents.application.services.acms_service import (  # noqa: E402
    DefaultStrategySelector,
    RecencyStrategy,
    RelevanceStrategy,
    StrategyCapabilities,
    TieredStrategy,
)
from cleveragents.domain.models.core.context_fragment import (  # noqa: E402
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)

_DEFAULT_PROV = FragmentProvenance(resource_uri="test://robot-orch")


def _make_frag(**kwargs: Any) -> ContextFragment:
    kwargs.setdefault("uko_node", "test://robot-orch")
    kwargs.setdefault("token_count", 0)
    kwargs.setdefault("provenance", _DEFAULT_PROV)
    return ContextFragment(**kwargs)


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def _cmd_cws_select() -> int:
    """ConfidenceWeightedSelector selects strategies."""
    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    selector = ConfidenceWeightedSelector()
    results = selector.select(strategies, {})
    assert len(results) == 3, f"Expected 3, got {len(results)}"
    confs = [c for _, c in results]
    assert confs == sorted(confs, reverse=True), "Not sorted"
    print("cws-select-ok")
    return 0


def _cmd_cws_preferred() -> int:
    """ConfidenceWeightedSelector boosts preferred strategies."""
    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    selector = ConfidenceWeightedSelector(preference_boost=1.5)
    baseline = dict(selector.select(strategies, {}))
    boosted = dict(selector.select(strategies, {"preferred_strategies": ["recency"]}))
    # recency should be boosted
    recency_baseline = baseline.get(
        next(s for s in strategies if s.name == "recency"), 0.0
    )
    recency_boosted = boosted.get(
        next(s for s in strategies if s.name == "recency"), 0.0
    )
    assert recency_boosted >= recency_baseline
    assert recency_boosted <= 1.0
    print("cws-preferred-ok")
    return 0


def _cmd_pba_proportional() -> int:
    """ProportionalBudgetAllocator distributes proportionally."""
    s1 = RelevanceStrategy()
    s2 = RecencyStrategy()
    allocator = ProportionalBudgetAllocator()
    allocs = allocator.allocate([(s1, 0.8), (s2, 0.2)], 1000)
    total = sum(a[2] for a in allocs)
    assert total == 1000, f"Total {total} != 1000"
    assert abs(allocs[0][2] - 800) <= 1
    assert abs(allocs[1][2] - 200) <= 1
    print("pba-proportional-ok")
    return 0


def _cmd_pba_min_budget() -> int:
    """ProportionalBudgetAllocator enforces min_useful_budget."""
    s1 = RelevanceStrategy()
    s2 = RecencyStrategy()
    allocator = ProportionalBudgetAllocator(min_useful_budget=200)
    allocs = allocator.allocate([(s1, 0.99), (s2, 0.01)], 300)
    assert len(allocs) == 1, f"Expected 1, got {len(allocs)}"
    assert allocs[0][2] == 300
    print("pba-min-budget-ok")
    return 0


def _cmd_cb_open() -> int:
    """CircuitBreaker opens after threshold failures."""
    cb = CircuitBreaker(failure_threshold=3)
    for _ in range(3):
        cb.record_failure("flaky")
    assert cb.is_open("flaky"), "Circuit should be open"
    print("cb-open-ok")
    return 0


def _cmd_cb_reset() -> int:
    """CircuitBreaker resets on success."""
    cb = CircuitBreaker(failure_threshold=3)
    cb.record_failure("recovering")
    cb.record_failure("recovering")
    cb.record_success("recovering")
    assert not cb.is_open("recovering"), "Circuit should be closed"
    print("cb-reset-ok")
    return 0


class _TrackingRobotStrategy:
    """Test strategy that tracks invocation."""

    def __init__(self) -> None:
        self.invoked = False

    @property
    def name(self) -> str:
        return "tracking_robot"

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities()

    def can_handle(self, request: dict[str, Any]) -> float:
        return 0.7

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        self.invoked = True
        result: list[ContextFragment] = []
        total = 0
        for frag in fragments:
            if total + frag.token_count <= budget.available_tokens:
                result.append(frag)
                total += frag.token_count
        return result

    def explain(self) -> str:
        return "Tracking."


def _cmd_executor_run() -> int:
    """ParallelStrategyExecutor invokes strategy."""
    strategy = _TrackingRobotStrategy()
    frags = [
        _make_frag(
            uko_node="project://t/a.py",
            content="alpha",
            token_count=50,
            relevance_score=0.9,
        ),
    ]
    executor = ParallelStrategyExecutor()
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    results = executor.execute([(strategy, 0.7, 200)], frags, budget)
    assert len(results) > 0, "Expected fragments"
    assert strategy.invoked, "Strategy not invoked"
    print("executor-run-ok")
    return 0


def _cmd_pipeline_assemble() -> int:
    """ContextAssemblyPipeline assembly with timing."""
    pipeline = ContextAssemblyPipeline()
    frags = [
        _make_frag(
            uko_node="project://t/a.py",
            content="alpha",
            token_count=100,
            relevance_score=0.9,
        ),
        _make_frag(
            uko_node="project://t/b.py",
            content="beta",
            token_count=100,
            relevance_score=0.3,
        ),
    ]
    budget = ContextBudget(max_tokens=250, reserved_tokens=0)
    payload = pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="relevance",
    )
    assert len(payload.fragments) == 2
    timings = pipeline.last_timings
    assert timings is not None
    assert timings.total_ms >= 0.0
    for name in type(timings).model_fields:
        assert getattr(timings, name) >= 0.0
    print("pipeline-assemble-ok")
    return 0


class _TrackingSelector:
    """Tracks invocation."""

    def __init__(self) -> None:
        self.called = False

    def select(
        self,
        strategies: Sequence[Any],
        request: dict[str, Any],
    ) -> list[tuple[Any, float]]:
        self.called = True
        return DefaultStrategySelector().select(strategies, request)


def _cmd_pipeline_custom() -> int:
    """ContextAssemblyPipeline with custom components."""
    selector = _TrackingSelector()
    pipeline = ContextAssemblyPipeline(strategy_selector=selector)
    frags = [
        _make_frag(
            uko_node="project://t/a.py",
            content="alpha",
            token_count=100,
            relevance_score=0.9,
        ),
    ]
    budget = ContextBudget(max_tokens=500, reserved_tokens=0)
    pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="relevance",
    )
    assert selector.called, "Custom selector not invoked"
    print("pipeline-custom-ok")
    return 0


# ---------------------------------------------------------------------------
# CLI dispatcher
# ---------------------------------------------------------------------------

_COMMANDS: dict[str, Callable[[], int]] = {
    "cws-select": _cmd_cws_select,
    "cws-preferred": _cmd_cws_preferred,
    "pba-proportional": _cmd_pba_proportional,
    "pba-min-budget": _cmd_pba_min_budget,
    "cb-open": _cmd_cb_open,
    "cb-reset": _cmd_cb_reset,
    "executor-run": _cmd_executor_run,
    "pipeline-assemble": _cmd_pipeline_assemble,
    "pipeline-custom": _cmd_pipeline_custom,
}


def main() -> int:
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>", file=sys.stderr)
        return 1
    try:
        return _COMMANDS[sys.argv[1]]()
    except Exception as exc:
        print(f"FAIL: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())
