"""Tests for alive.core.composer."""

import asyncio
import json
import sys
from pathlib import Path

import pytest
from rich.console import Console

from alive.core.composer import (
    CompositionError,
    CompositionResult,
    Composer,
    StepResult,
    compose,
)
from alive.core.planner import ExecutionPlan, PlanStep
from alive.core.skill_registry import SkillRegistry
from alive.llm.base import LLMProvider
from alive.sandbox.dependency import DependencyManager


# ── Fixtures ─────────────────────────────────────────────────────────────────


class MockProvider(LLMProvider):
    """LLM provider that returns canned responses."""

    name = "mock"

    def __init__(self, responses: list[str] | None = None):
        self._responses = list(responses or [])
        self._call_count = 0
        self.prompts: list[str] = []

    async def complete(self, system_prompt: str, user_prompt: str) -> str:
        self.prompts.append(user_prompt)
        if self._call_count < len(self._responses):
            resp = self._responses[self._call_count]
        else:
            resp = self._responses[-1] if self._responses else ""
        self._call_count += 1
        return resp


class MockDependencyManager(DependencyManager):
    """Uses system Python — no real venv needed."""

    def __init__(self):
        self._venv_dir = Path(sys.executable).parent.parent

    @property
    def python(self) -> Path:
        return Path(sys.executable)

    def ensure_venv(self) -> None:
        pass

    def install(self, packages: list[str]) -> None:
        pass


def _abort_handler(step_index, step, error, console):
    return "abort"


def _skip_handler(step_index, step, error, console):
    return "skip"


def _retry_then_abort_handler():
    """Returns retry on first call, abort on second."""
    calls = []

    def handler(step_index, step, error, console):
        calls.append(1)
        return "retry" if len(calls) == 1 else "abort"

    return handler


@pytest.fixture
def registry(tmp_path):
    """Create an isolated registry for testing."""
    import alive.seeds._installer as installer

    original = installer._SEEDS_DIR
    installer._SEEDS_DIR = tmp_path / "empty_seeds"
    installer._SEEDS_DIR.mkdir()
    reg = SkillRegistry(skills_dir=tmp_path / "skills")
    installer._SEEDS_DIR = original
    return reg


@pytest.fixture
def mock_dm():
    return MockDependencyManager()


@pytest.fixture
def quiet_console():
    return Console(quiet=True)


def _make_skill(registry, tmp_path_for_skill, name="adder", code=None):
    """Register a simple skill and write its file for sandbox execution."""
    code = code or 'def run(x: int = 1, y: int = 2) -> dict:\n    return {"sum": x + y, "success": True}\n'
    registry.register(name=name, code=code, description=f"Test skill {name}")
    return name


# ── Data structure tests ─────────────────────────────────────────────────────


def test_step_result_defaults():
    step = PlanStep(type="LLM_ONLY", description="test")
    sr = StepResult(step_index=0, step=step, success=True, output="hi")
    assert sr.error is None
    assert sr.duration_ms == 0
    assert sr.skipped is False


def test_composition_result_creation():
    plan = ExecutionPlan(steps=[], reasoning="", original_request="test")
    cr = CompositionResult(
        plan=plan,
        step_results=[],
        final_output="done",
        success=True,
    )
    assert cr.success is True
    assert cr.final_output == "done"


# ── LLM_ONLY step tests ─────────────────────────────────────────────────────


def test_llm_only_single_step(registry, quiet_console):
    """Single LLM_ONLY step should execute and return the LLM response."""
    plan = ExecutionPlan(
        steps=[PlanStep(type="LLM_ONLY", description="answer", prompt="What is 2+2?")],
        reasoning="simple question",
        original_request="what is 2+2",
    )
    provider = MockProvider(["4"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    assert result.success is True
    assert len(result.step_results) == 1
    assert result.step_results[0].success is True
    # Single step — no synthesis, output returned directly
    assert result.final_output == "4"


def test_llm_only_with_context(registry, quiet_console):
    """Second LLM_ONLY step should receive context from the first."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="step 1", prompt="Say hello"),
            PlanStep(type="LLM_ONLY", description="step 2", prompt="Now say goodbye"),
        ],
        reasoning="two steps",
        original_request="greet and farewell",
    )
    provider = MockProvider(["Hello!", "Goodbye!", "Final synthesis"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    assert result.success is True
    assert len(result.step_results) == 2

    # The second prompt should contain context from step 1
    second_prompt = provider.prompts[1]
    assert "Hello!" in second_prompt
    assert "Previous step results" in second_prompt


# ── EXISTING_SKILL step tests ───────────────────────────────────────────────


def test_existing_skill_success(registry, mock_dm, quiet_console, monkeypatch):
    """EXISTING_SKILL step should run the skill and return output."""
    _make_skill(registry, None, name="adder")

    plan = ExecutionPlan(
        steps=[
            PlanStep(
                type="EXISTING_SKILL",
                description="add numbers",
                skill="adder",
                input={"x": 3, "y": 5},
            ),
        ],
        reasoning="direct",
        original_request="add 3 and 5",
    )

    # Mock run_skill to avoid real subprocess
    import alive.core.composer as comp_mod

    def mock_run_skill(name, input_data=None, **kwargs):
        from alive.sandbox.runner import SkillResult

        return SkillResult(
            success=True,
            output={"sum": 8, "success": True},
            error=None,
            duration_ms=10,
        )

    monkeypatch.setattr(comp_mod, "run_skill", mock_run_skill, raising=False)

    # Patch run_skill import inside the method
    original_run_existing = Composer._run_existing_skill

    def patched_run_existing(self, index, step, reg):
        from alive.sandbox.runner import SkillResult as SR

        sr = mock_run_skill(step.skill, step.input)
        try:
            reg.update_stats(step.skill or "", success=sr.success)
        except KeyError:
            pass
        return StepResult(
            step_index=index,
            step=step,
            success=True,
            output=json.dumps(sr.output, default=str),
        )

    monkeypatch.setattr(Composer, "_run_existing_skill", patched_run_existing)

    provider = MockProvider(["synthesis"])
    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    assert result.success is True
    assert '"sum": 8' in result.step_results[0].output


def test_existing_skill_failure(registry, quiet_console, monkeypatch):
    """Failed EXISTING_SKILL should report error and trigger error handler."""
    _make_skill(registry, None, name="broken")

    plan = ExecutionPlan(
        steps=[
            PlanStep(
                type="EXISTING_SKILL",
                description="run broken",
                skill="broken",
                input={},
            ),
        ],
        reasoning="will fail",
        original_request="break",
    )

    def mock_run_skill(name, input_data=None, **kwargs):
        from alive.sandbox.runner import SkillResult

        return SkillResult(
            success=False,
            output=None,
            error="ValueError: boom",
            duration_ms=5,
        )

    monkeypatch.setattr(Composer, "_run_existing_skill", lambda self, i, s, r: StepResult(
        step_index=i, step=s, success=False, output="", error="ValueError: boom",
    ))

    provider = MockProvider([])

    with pytest.raises(CompositionError):
        asyncio.run(
            compose(
                plan,
                provider=provider,
                registry=registry,
                console=quiet_console,
                on_error=_abort_handler,
            )
        )


# ── Error handling tests ─────────────────────────────────────────────────────


def test_skip_failed_step(registry, quiet_console, monkeypatch):
    """Skipping a failed step should continue execution."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="fail step", prompt="fail"),
            PlanStep(type="LLM_ONLY", description="ok step", prompt="succeed"),
        ],
        reasoning="test skip",
        original_request="test",
    )

    call_count = [0]
    original_run_llm = Composer._run_llm_only

    async def failing_first(self, index, step, provider):
        call_count[0] += 1
        if index == 0:
            return StepResult(
                step_index=index, step=step, success=False, output="", error="forced failure",
            )
        return await original_run_llm(self, index, step, provider)

    monkeypatch.setattr(Composer, "_run_llm_only", failing_first)

    provider = MockProvider(["ok result", "synthesis"])
    result = asyncio.run(
        compose(
            plan,
            provider=provider,
            registry=registry,
            console=quiet_console,
            on_error=_skip_handler,
        )
    )
    assert result.step_results[0].skipped is True
    assert result.step_results[1].success is True


def test_abort_raises_composition_error(registry, quiet_console, monkeypatch):
    """Aborting should raise CompositionError."""
    plan = ExecutionPlan(
        steps=[PlanStep(type="LLM_ONLY", description="fail", prompt="fail")],
        reasoning="",
        original_request="test",
    )

    async def always_fail(self, index, step, provider):
        return StepResult(
            step_index=index, step=step, success=False, output="", error="nope",
        )

    monkeypatch.setattr(Composer, "_run_llm_only", always_fail)

    with pytest.raises(CompositionError, match="Step 1"):
        asyncio.run(
            compose(
                plan,
                provider=MockProvider([]),
                registry=registry,
                console=quiet_console,
                on_error=_abort_handler,
            )
        )


def test_retry_then_succeed(registry, quiet_console, monkeypatch):
    """Retry should re-execute the step."""
    plan = ExecutionPlan(
        steps=[PlanStep(type="LLM_ONLY", description="flaky", prompt="work")],
        reasoning="",
        original_request="test",
    )

    attempts = [0]

    async def flaky_step(self, index, step, provider):
        attempts[0] += 1
        if attempts[0] == 1:
            return StepResult(
                step_index=index, step=step, success=False, output="", error="first try failed",
            )
        return StepResult(
            step_index=index, step=step, success=True, output="worked on retry",
        )

    monkeypatch.setattr(Composer, "_run_llm_only", flaky_step)

    handler = _retry_then_abort_handler()
    result = asyncio.run(
        compose(
            plan,
            provider=MockProvider([]),
            registry=registry,
            console=quiet_console,
            on_error=handler,
        )
    )
    assert result.success is True
    assert attempts[0] == 2
    assert result.step_results[0].output == "worked on retry"


# ── Multi-step chaining tests ───────────────────────────────────────────────


def test_multi_step_chaining(registry, quiet_console):
    """Three LLM_ONLY steps should chain context correctly."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="first", prompt="Write a greeting"),
            PlanStep(type="LLM_ONLY", description="second", prompt="Translate it"),
            PlanStep(type="LLM_ONLY", description="third", prompt="Make it formal"),
        ],
        reasoning="chained",
        original_request="formal greeting in French",
    )
    provider = MockProvider(["Bonjour!", "Hello!", "Dear Sir", "Final summary"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    assert result.success is True
    assert len(result.step_results) == 3

    # Third prompt should have context from first two steps
    third_prompt = provider.prompts[2]
    assert "Bonjour!" in third_prompt
    assert "Hello!" in third_prompt


# ── Synthesis tests ──────────────────────────────────────────────────────────


def test_single_step_skips_synthesis(registry, quiet_console):
    """Single successful step should return its output directly without synthesis."""
    plan = ExecutionPlan(
        steps=[PlanStep(type="LLM_ONLY", description="just one", prompt="answer")],
        reasoning="",
        original_request="test",
    )
    provider = MockProvider(["direct answer"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    # Should NOT have made a synthesis call (only 1 LLM call total)
    assert provider._call_count == 1
    assert result.final_output == "direct answer"


def test_multi_step_synthesis(registry, quiet_console):
    """Multiple steps should trigger a synthesis LLM call."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="step 1", prompt="first"),
            PlanStep(type="LLM_ONLY", description="step 2", prompt="second"),
        ],
        reasoning="",
        original_request="test",
    )
    provider = MockProvider(["result 1", "result 2", "synthesized final"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    # 2 step calls + 1 synthesis = 3 total
    assert provider._call_count == 3
    assert result.final_output == "synthesized final"

    # Synthesis prompt should contain both results
    synth_prompt = provider.prompts[2]
    assert "result 1" in synth_prompt
    assert "result 2" in synth_prompt


# ── Display test ─────────────────────────────────────────────────────────────


def test_execute_with_quiet_console_no_crash(registry, quiet_console):
    """Full execution with quiet console should not raise."""
    plan = ExecutionPlan(
        steps=[PlanStep(type="LLM_ONLY", description="ok", prompt="hello")],
        reasoning="test",
        original_request="hi",
    )
    provider = MockProvider(["world"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    assert result.success is True


# ── CompositionError tests ───────────────────────────────────────────────────


def test_composition_error_message():
    step = PlanStep(type="LLM_ONLY", description="broken step")
    err = CompositionError(2, step, "something broke")
    assert "Step 3" in str(err)
    assert "broken step" in str(err)
    assert "something broke" in str(err)
    assert err.step_index == 2


# ── Additional composer tests ───────────────────────────────────────────────


def test_no_successful_steps_returns_message(registry, quiet_console, monkeypatch):
    """When all steps are skipped, synthesis says no steps completed."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="step 1", prompt="x"),
            PlanStep(type="LLM_ONLY", description="step 2", prompt="y"),
        ],
        reasoning="",
        original_request="test",
    )

    async def always_fail(self, index, step, provider):
        return StepResult(
            step_index=index, step=step, success=False, output="", error="fail",
        )

    monkeypatch.setattr(Composer, "_run_llm_only", always_fail)

    result = asyncio.run(
        compose(
            plan,
            provider=MockProvider([]),
            registry=registry,
            console=quiet_console,
            on_error=_skip_handler,
        )
    )
    assert "No steps completed" in result.final_output


def test_synthesis_fallback_on_provider_error(registry, quiet_console, monkeypatch):
    """When synthesis LLM call fails, fallback to concatenated output."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="step 1", prompt="first"),
            PlanStep(type="LLM_ONLY", description="step 2", prompt="second"),
        ],
        reasoning="",
        original_request="test",
    )

    call_count = [0]

    class FailingSynthProvider(LLMProvider):
        name = "fail_synth"

        async def complete(self, system_prompt: str, user_prompt: str) -> str:
            call_count[0] += 1
            if call_count[0] <= 2:
                return f"output {call_count[0]}"
            raise RuntimeError("synthesis failed")

    result = asyncio.run(
        compose(
            plan,
            provider=FailingSynthProvider(),
            registry=registry,
            console=quiet_console,
        )
    )
    # Should fallback to concatenated outputs
    assert "step 1" in result.final_output
    assert "step 2" in result.final_output


def test_llm_only_step_failure_on_exception(registry, quiet_console):
    """LLM provider exception in LLM_ONLY step should return failure, not crash."""
    plan = ExecutionPlan(
        steps=[PlanStep(type="LLM_ONLY", description="will error", prompt="crash")],
        reasoning="",
        original_request="test",
    )

    class ErrorProvider(LLMProvider):
        name = "error"
        async def complete(self, system_prompt: str, user_prompt: str) -> str:
            raise ConnectionError("network down")

    with pytest.raises(CompositionError):
        asyncio.run(
            compose(
                plan,
                provider=ErrorProvider(),
                registry=registry,
                console=quiet_console,
                on_error=_abort_handler,
            )
        )


def test_build_context_truncation(registry, quiet_console):
    """Long step outputs should be truncated in context passed to next step."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="verbose", prompt="long output"),
            PlanStep(type="LLM_ONLY", description="consumer", prompt="use it"),
        ],
        reasoning="",
        original_request="test",
    )
    # First response is over 2000 chars
    long_output = "x" * 3000
    provider = MockProvider([long_output, "consumed", "synthesis"])

    result = asyncio.run(
        compose(plan, provider=provider, registry=registry, console=quiet_console)
    )
    # Second prompt should have truncated context
    second_prompt = provider.prompts[1]
    # Context is truncated to 2000 chars
    assert len(second_prompt) < 3000 + 500  # some overhead for template text


def test_build_context_includes_skipped(registry, quiet_console, monkeypatch):
    """Context should show [skipped] for skipped steps."""
    plan = ExecutionPlan(
        steps=[
            PlanStep(type="LLM_ONLY", description="skip me", prompt="x"),
            PlanStep(type="LLM_ONLY", description="use context", prompt="y"),
        ],
        reasoning="",
        original_request="test",
    )

    call_count = [0]
    original = Composer._run_llm_only

    async def fail_first(self, index, step, provider):
        call_count[0] += 1
        if index == 0:
            return StepResult(
                step_index=0, step=step, success=False, output="", error="oops",
            )
        return await original(self, index, step, provider)

    monkeypatch.setattr(Composer, "_run_llm_only", fail_first)

    provider = MockProvider(["ok result", "synthesis"])
    result = asyncio.run(
        compose(
            plan,
            provider=provider,
            registry=registry,
            console=quiet_console,
            on_error=_skip_handler,
        )
    )
    # The second step's prompt should contain "[skipped]"
    second_prompt = provider.prompts[0]  # first actual LLM call after skip
    # Check context in second step's prompt or synthesis
    assert result.step_results[0].skipped is True


def test_interactive_error_handler_eof_returns_abort(quiet_console, monkeypatch):
    """_interactive_error_handler should return 'abort' on EOFError."""
    from alive.core.composer import _interactive_error_handler

    step = PlanStep(type="LLM_ONLY", description="test")
    monkeypatch.setattr("alive.core.composer.Prompt.ask", lambda *a, **kw: (_ for _ in ()).throw(EOFError))

    action = _interactive_error_handler(0, step, "test error", quiet_console)
    assert action == "abort"


def test_composition_error_attributes():
    """CompositionError should store step_index and step."""
    step = PlanStep(type="NEW_SKILL", description="build widget", capability="widgets")
    err = CompositionError(0, step, "build failed")
    assert err.step_index == 0
    assert err.step == step
    assert isinstance(err, RuntimeError)
