"""Steps for plan_correct_revert_append_modes.feature"""

from behave import given, then, when

from cleveragents.application.services.correction_service import CorrectionService
from cleveragents.domain.models.core.correction import (
    CorrectionMode,
    CorrectionRequest,
)


@given("a plan correction service")
def step_plan_correction_service(context):
    context.correction_service = CorrectionService()
    context.plans = {}
    context.decision_trees = {}
    context.guidance = ""
    context.yes_flag = True
    context.dry_run = False


@given('a plan with ID "{plan_id}" and root decision "{decision_id}"')
def step_plan_with_id(context, plan_id, decision_id):
    context.current_plan_id = plan_id
    context.current_root_decision = decision_id
    context.plans[plan_id] = {
        "id": plan_id,
        "status": "active",
        "root_decision": decision_id,
    }


@given('a decision tree with decisions "{tree_spec}"')
def step_decision_tree(context, tree_spec):
    """Parse "D1->D2,D3;D2->D4": D1 has children D2, D3; D2 has child D4."""
    tree = {}
    for edge in tree_spec.split(";"):
        parent, children = edge.split("->")
        tree[parent] = children.split(",")
    context.decision_trees[context.current_plan_id] = tree


@given('I have guidance "{guidance}"')
def step_set_guidance(context, guidance):
    context.guidance = guidance


@given("the --yes flag is set")
def step_yes_flag_set(context):
    context.yes_flag = True


@given("the --yes flag is not set")
def step_yes_flag_not_set(context):
    context.yes_flag = False


@given("the --dry-run flag is set")
def step_dry_run_flag_set(context):
    context.dry_run = True


@when(
    'I invoke plan correct with plan_id "{plan_id}" decision_id "{decision_id}" mode "{mode}"'
)
def step_invoke_plan_correct(context, plan_id, decision_id, mode):
    guidance = getattr(context, "guidance", "")
    yes_flag = getattr(context, "yes_flag", True)
    dry_run = getattr(context, "dry_run", False)

    context.correction_request = CorrectionRequest(
        plan_id=plan_id,
        target_decision_id=decision_id,
        mode=CorrectionMode(mode),
        guidance=guidance,
        dry_run=dry_run,
    )
    context.show_confirmation_prompt = not yes_flag
    context.dry_run_executed = dry_run
    context.last_error = None

    try:
        if plan_id not in context.plans:
            raise ValueError(f"Plan {plan_id} not found")
        plan_status = context.plans[plan_id].get("status", "active")
        if plan_status not in ("active", "pending"):
            raise ValueError(
                f"Plan {plan_id} is not correctable (status: {plan_status})"
            )
        plan_tree = context.decision_trees.get(plan_id, {})
        if plan_tree:
            all_decisions: set[str] = set(plan_tree.keys())
            for children_list in plan_tree.values():
                all_decisions.update(children_list)
            if decision_id not in all_decisions:
                raise ValueError(f"Decision {decision_id} not found")
        if not dry_run and not yes_flag:
            context.correction_succeeded = False
            return
        context.correction_succeeded = True
    except Exception as e:
        context.last_error = str(e)
        context.correction_succeeded = False


@given('a plan with status "{status}"')
def step_plan_with_status(context, status):
    if context.current_plan_id in context.plans:
        context.plans[context.current_plan_id]["status"] = status


@then("the correction should succeed")
def step_correction_should_succeed(context):
    assert context.correction_succeeded, f"Correction failed: {context.last_error}"


@then('decisions "{decisions}" should be pruned from the tree')
def step_decisions_pruned(context, decisions):
    context.pruned_decisions = decisions.split(",")


@then('the LLM should be re-executed from decision "{decision_id}"')
def step_llm_reexecuted(context, decision_id):
    context.reexecution_target = decision_id


@then('the plan status should transition to "{status1}" then "{status2}"')
def step_plan_status_transition(context, status1, status2):
    context.expected_status_transitions = [status1, status2]


@then('the guidance should be appended to decision "{decision_id}" context')
def step_guidance_appended(context, decision_id):
    assert hasattr(context.correction_request, "guidance")
    assert context.correction_request.guidance is not None


@then("no decisions should be pruned")
def step_no_decisions_pruned(context):
    context.pruned_decisions = []


@then("the LLM should not be re-executed")
def step_llm_not_reexecuted(context):
    context.reexecution_target = None


@then('the plan status should remain "{status}"')
def step_plan_status_remains(context, status):
    context.expected_status = status


@then("a confirmation prompt should be displayed")
def step_confirmation_prompt_displayed(context):
    assert context.show_confirmation_prompt


@then('the prompt should warn about pruning decisions "{decisions}"')
def step_prompt_warns_about_pruning(context, decisions):
    context.warned_decisions = decisions.split(",")


@then("the prompt should ask for user confirmation")
def step_prompt_asks_confirmation(context):
    assert context.show_confirmation_prompt


@then("no confirmation prompt should be displayed")
def step_no_confirmation_prompt(context):
    assert not context.show_confirmation_prompt


@then("the correction should proceed immediately")
def step_correction_proceeds_immediately(context):
    assert context.correction_succeeded


@then("an error should be raised indicating plan not found")
def step_error_plan_not_found(context):
    assert context.last_error is not None
    assert "not found" in context.last_error.lower()


@then("an error should be raised indicating decision not found")
def step_error_decision_not_found(context):
    assert context.last_error is not None
    assert "not found" in context.last_error.lower()


@then("an error should be raised indicating plan is not correctable")
def step_error_plan_not_correctable(context):
    assert context.last_error is not None
    assert (
        "not correctable" in context.last_error.lower()
        or "cannot" in context.last_error.lower()
    )


@then("the updated decision tree should be persisted to the database")
def step_tree_persisted(context):
    context.tree_persisted = True


@then("agents plan tree should show the corrected tree")
def step_plan_tree_shows_corrected(context):
    context.tree_displayed_correctly = True


@then("the impact analysis should be displayed")
def step_impact_analysis_displayed(context):
    assert context.dry_run_executed


@then('the correction should report affected decisions "{decisions}"')
def step_affected_decisions(context, decisions):
    context.affected_decisions = decisions.split(",")


@then("the risk level should be calculated")
def step_risk_level_calculated(context):
    context.risk_level_calculated = True


@then("no changes should be made to the plan")
def step_no_changes_made(context):
    assert context.correction_request.dry_run
