from __future__ import annotations

import asyncio
import json
from copy import deepcopy
from typing import Any

from behave import given, then, when

from cleveragents.agents.graphs.auto_debug import AutoDebugAgent, AutoDebugState


class _SequentialLLM:
    """Stub LLM that returns sequential responses with a content attribute."""

    def __init__(self, responses: list[str]):
        self._responses = list(responses)
        self.invocations: list[list[Any]] = []

    def invoke(self, messages):
        self.invocations.append(messages)
        if not self._responses:
            raise AssertionError("No responses left in stub LLM")
        content = self._responses.pop(0)
        return type("Response", (), {"content": content})()


@given("a custom auto debug LLM yields JSON responses")
def step_custom_llm(context):
    context.custom_llm = _SequentialLLM(
        [
            "Custom analysis output",
            json.dumps(
                {
                    "description": "JSON described fix",
                    "code": "# patched code",
                    "files_to_modify": ["file.py"],
                }
            ),
            json.dumps(
                {
                    "is_valid": True,
                    "reasoning": "Structured JSON shows success",
                    "issues": [],
                }
            ),
        ]
    )


@when("I run the auto debug agent graph with sample state")
def step_run_auto_debug_graph(context):
    context.agent = AutoDebugAgent(llm=context.custom_llm, max_fix_attempts=2)

    input_state: AutoDebugState = {
        "messages": [],
        "context": {},
        "result": None,
        "error": None,
        "metadata": {},
        "error_message": "NameError: missing value",
        "code_context": "value = missing",
        "attempted_fixes": [],
        "current_fix": {},
        "fix_validated": False,
    }

    context.final_state = context.agent.invoke(input_state)


@then("the provided LLM instance should be used")
def step_assert_llm_used(context):
    assert context.agent.llm is context.custom_llm
    assert len(context.custom_llm.invocations) == 3


@then("the fix JSON should populate the current fix")
def step_assert_fix_parsed(context):
    fix = context.final_state["current_fix"]
    assert fix["description"] == "JSON described fix"
    assert fix["code"] == "# patched code"
    assert fix["files_to_modify"] == ["file.py"]


@then("the validation JSON should mark the fix as valid")
def step_assert_validation_parsed(context):
    assert context.final_state["fix_validated"] is True
    result = context.final_state.get("result")
    assert result is not None
    assert result["success"] is True
    assert result["fix"] == context.final_state["current_fix"]
    assert result["attempts"] == 0


@given("a retrying auto debug LLM yields a failed then successful validation")
def step_retrying_llm(context):
    context.retry_llm = _SequentialLLM(
        [
            "Initial analysis output",
            json.dumps(
                {
                    "description": "First fix",
                    "code": "# first attempt",
                    "files_to_modify": ["first.py"],
                }
            ),
            json.dumps(
                {
                    "is_valid": False,
                    "reasoning": "Validation failed",
                    "issues": ["still broken"],
                }
            ),
            json.dumps(
                {
                    "description": "Second fix",
                    "code": "# second attempt",
                    "files_to_modify": ["second.py"],
                }
            ),
            json.dumps(
                {
                    "is_valid": True,
                    "reasoning": "Validation succeeded",
                    "issues": [],
                }
            ),
        ]
    )


@when("I run the auto debug agent graph with a retry")
def step_run_auto_debug_with_retry(context):
    agent = AutoDebugAgent(llm=context.retry_llm, max_fix_attempts=2)

    input_state: AutoDebugState = {
        "messages": [],
        "context": {},
        "result": None,
        "error": None,
        "metadata": {},
        "error_message": "NameError: missing value",
        "code_context": "value = missing",
        "attempted_fixes": [],
        "current_fix": {},
        "fix_validated": False,
    }

    context.retry_state = agent.invoke(input_state)


@then("the agent should retry and capture the failed fix")
def step_assert_retry_behavior(context):
    state = context.retry_state
    assert state["fix_validated"] is True
    assert len(state["attempted_fixes"]) == 1
    assert state["attempted_fixes"][0]["description"] == "First fix"
    assert state["current_fix"]["description"] == "Second fix"
    result = state.get("result")
    assert result is not None
    assert result["success"] is True
    assert result["attempts"] == 1


@given("a simple auto debug LLM for async calls")
def step_async_llm(context):
    context.async_llm_responses = [
        "Async analysis output",
        json.dumps(
            {
                "description": "Async fix",
                "code": "# async patched",
                "files_to_modify": ["async.py"],
            }
        ),
        json.dumps(
            {
                "is_valid": True,
                "reasoning": "Async validation succeeded",
                "issues": [],
            }
        ),
    ]


@when("I run the auto debug async interfaces")
def step_run_async_interfaces(context):
    base_state: AutoDebugState = {
        "messages": [],
        "context": {},
        "result": None,
        "error": None,
        "metadata": {},
        "error_message": "NameError: missing value",
        "code_context": "value = missing",
        "attempted_fixes": [],
        "current_fix": {},
        "fix_validated": False,
    }

    def make_agent() -> AutoDebugAgent:
        return AutoDebugAgent(
            llm=_SequentialLLM(list(context.async_llm_responses)), max_fix_attempts=1
        )

    async def collect_async_results():
        agent_for_ainvoke = make_agent()
        context.async_invoke_state = await agent_for_ainvoke.ainvoke(
            deepcopy(base_state)
        )

        agent_for_stream = make_agent()
        context.stream_events = list(agent_for_stream.stream(deepcopy(base_state)))

        agent_for_astream = make_agent()
        async_events: list[dict[str, Any]] = []
        async for event in agent_for_astream.astream(deepcopy(base_state)):
            async_events.append(event)
        context.async_stream_events = async_events

    asyncio.run(collect_async_results())


@then("the async auto debug results should complete")
def step_assert_async_results(context):
    assert context.async_invoke_state["result"]["success"] is True

    assert context.stream_events
    assert isinstance(context.stream_events[-1], dict)

    assert context.async_stream_events
    assert isinstance(context.async_stream_events[-1], dict)

    def extract_success(events):
        for event in events:
            if not isinstance(event, dict):
                continue
            for value in event.values():
                if not isinstance(value, dict):
                    continue
                result = value.get("result")
                if isinstance(result, dict) and "success" in result:
                    return result.get("success")
        return None

    stream_success = extract_success(context.stream_events)
    if stream_success is not None:
        assert stream_success is True

    async_success = extract_success(context.async_stream_events)
    if async_success is not None:
        assert async_success is True