"""The correction loop: the gate hands back structured errors for an invalid capability manifest,
and Core re-dispatches the SAME run so the agent fixes it on top of its own saved work (bounded),
instead of ending the run or regenerating the frontend."""

from __future__ import annotations

import pickle
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch
from uuid import uuid4

import pytest
from src.config import get_settings
from src.domain.codegen.provision_gate import PrimitiveGapError, request_provisioning
from src.domain.codegen.schemas import CodegenActionsCompleteBody
from src.domain.codegen.service import CodegenService
from src.domain.codegen.work_branch import apply_report, new_state
from src.workers.tasks.codegen import _correction_from_outcomes

BASE = "a" * 40
CKPT = "b" * 40
DETAILS = [{"path": "fields.0.public", "code": "unknown_key", "message": "not recognized"}]


def settings(enabled=True, corrections=2):
    return get_settings().model_copy(
        update={"codegen_work_branch_enabled": enabled,
                "codegen_manifest_correction_attempts": corrections}
    )


def make_service(run):
    repo = SimpleNamespace(
        get_by_id=AsyncMock(return_value=run),
        update_status=AsyncMock(),
        update_checkpoint=AsyncMock(),
        commit=AsyncMock(),
        refresh_started_at=AsyncMock(),
    )
    service = CodegenService(codegen_repo=repo, project_facade=SimpleNamespace())
    return service, repo


def make_run(state, status="running"):
    run_id = uuid4()
    state = {**state, "branch": f"p2hs/codegen/{run_id}"}
    return SimpleNamespace(
        id=run_id, project_id=uuid4(), status=status, worker_checkpoint="github_actions_dispatched",
        checkpoint_payload={}, execution_target="cloud",
        plan={"system_prompt": "sys", "user_prompt": "u", "context": {"work_branch": state}},
    )


def blocked_state(attempts=0):
    correction = {"entries": [{"name": "contact_submission", "errors": DETAILS}]}
    state = apply_report(
        new_state(uuid4(), "preview", BASE), status="blocked", checkpoint="agent_output",
        checkpoint_sha=CKPT, failure_class="invalid_requirement", retryable=True,
        correction=correction,
    )
    return {**state, "correction_attempts": attempts}


def invalid(name="contact_submission", criticality="required", **extra):
    return {"name": name, "ok": False, "status": "invalid_manifest", "criticality": criticality,
            "error_details": DETAILS, "renamed": [{"path": "name", "from": "Contact", "to": name}],
            **extra}


# ---- the gate hands the agent structured errors --------------------------------------------------


def test_the_correction_is_only_built_when_every_blocker_is_something_the_agent_can_fix():
    built = _correction_from_outcomes([], [invalid()])
    assert built["entries"][0]["name"] == "contact_submission"
    assert built["entries"][0]["errors"][0]["path"] == "fields.0.public"
    assert built["entries"][0]["renamed"][0]["to"] == "contact_submission"
    gap = {"name": "expense", "ok": False, "status": "primitive_gap", "criticality": "required"}
    # a real capability gap cannot be fixed by rewriting a file: no attempt is wasted on it
    assert _correction_from_outcomes([], [invalid(), gap]) is None
    assert _correction_from_outcomes([{"pack_slug": "p", "ok": False}], [invalid()]) is None
    # an OPTIONAL invalid entry does not block, so there is nothing to correct
    assert _correction_from_outcomes([], [invalid(criticality="optional")]) is None
    assert _correction_from_outcomes([], []) is None


def test_the_correction_travels_on_the_exception_through_arq_pickling():
    error = PrimitiveGapError("Fix the manifest", correction={"entries": [{"name": "x"}]})
    restored = pickle.loads(pickle.dumps(error))
    assert str(restored) == "Fix the manifest"
    assert restored.correction == {"entries": [{"name": "x"}]}
    assert PrimitiveGapError("plain gap").correction is None


class FakeRedis:
    async def enqueue_job(self, *args, **kwargs):
        return None

    async def aclose(self):
        return None


async def gate_result(error, monkeypatch):
    async def get_pool():
        return FakeRedis()

    monkeypatch.setattr("src.workers.enqueue._get_redis_pool", get_pool)

    class Job:
        def __init__(self, *args):
            pass

        async def result_info(self):
            return SimpleNamespace(success=False, result=error)

    monkeypatch.setattr("src.domain.codegen.provision_gate.Job", Job)
    repo = SimpleNamespace(
        get_by_id=AsyncMock(return_value=SimpleNamespace(status="running", project_id=uuid4())),
        get_project_discovery=AsyncMock(return_value=None),
    )
    return await request_provisioning(repo, uuid4(), {}, [], [])


@pytest.mark.asyncio
async def test_an_invalid_manifest_is_reported_as_correctable_and_a_real_gap_keeps_its_old_shape(
    monkeypatch,
):
    correction = {"entries": [{"name": "x", "errors": DETAILS}]}
    corrected = await gate_result(PrimitiveGapError("bad manifest", correction=correction), monkeypatch)
    assert corrected["status"] == "blocked"
    assert corrected["failure_class"] == "invalid_requirement" and corrected["retryable"] is True
    assert corrected["correction"] == correction
    plain = await gate_result(PrimitiveGapError("needs approval.quorum"), monkeypatch)
    assert plain == {"status": "blocked", "message": "needs approval.quorum"}


# ---- orchestration -------------------------------------------------------------------------------


async def correct(run, *, enabled=True, corrections=2, dispatch=None):
    service, repo = make_service(run)
    dispatch = dispatch or AsyncMock(return_value={"work_branch": "b", "dispatch_commit_sha": "d"})
    with patch("src.config.get_settings", return_value=settings(enabled, corrections)), \
         patch("src.domain.codegen.actions_pipeline.dispatch_codegen_actions_run", dispatch):
        result = await service._correct_github_actions_run(run)
    return result, dispatch, repo


@pytest.mark.asyncio
async def test_an_invalid_manifest_redispatches_the_same_run_with_the_errors_on_top_of_its_work():
    run = make_run(blocked_state())
    result, dispatch, repo = await correct(run)
    assert result["status"] == "correcting" and result["correction_attempt"] == 1
    ctx = dispatch.await_args.kwargs["plan_context"]
    assert dispatch.await_args.kwargs["run_id"] == run.id  # same run, same branch
    assert ctx["resume"]["skip_agent"] is False  # the agent DOES run, on top of its own output
    assert ctx["correction"]["attempt"] == 1
    assert ctx["correction"]["entries"][0]["errors"][0]["path"] == "fields.0.public"
    assert ctx["work_branch"]["status"] == "active" and ctx["work_branch"]["correction_attempts"] == 1
    assert repo.update_checkpoint.await_args.args[1] == "github_actions_dispatched"


@pytest.mark.asyncio
@pytest.mark.parametrize("kwargs", [{"enabled": False}, {"corrections": 0}])
async def test_no_correction_when_disabled_or_out_of_attempts(kwargs):
    result, dispatch, _ = await correct(make_run(blocked_state()), **kwargs)
    assert result is None
    dispatch.assert_not_awaited()


@pytest.mark.asyncio
async def test_correction_stops_after_the_last_allowed_attempt():
    result, dispatch, _ = await correct(make_run(blocked_state(attempts=2)), corrections=2)
    assert result is None
    dispatch.assert_not_awaited()


@pytest.mark.asyncio
async def test_a_failed_dispatch_leaves_the_plan_as_it_was_so_the_run_can_still_be_blocked():
    run = make_run(blocked_state())
    before = run.plan
    result, _, repo = await correct(run, dispatch=AsyncMock(side_effect=RuntimeError("github down")))
    assert result is None
    assert run.plan is before  # not left half-updated
    repo.update_checkpoint.assert_not_awaited()


def active_state():
    return new_state(uuid4(), "preview", BASE)


async def runner_reports_blocked(monkeypatch, corrected):
    run = make_run(active_state())
    service, repo = make_service(run)
    service._correct_github_actions_run = AsyncMock(return_value=corrected)
    monkeypatch.setattr("src.core.pubsub.get_pubsub_redis", AsyncMock(return_value=AsyncMock()))
    monkeypatch.setattr("src.core.run_events.publish_project_event_dual", AsyncMock())
    result = await service.report_run_complete(
        run.id,
        CodegenActionsCompleteBody(
            status="blocked", summary="invalid manifest", checkpoint="agent_output",
            checkpoint_sha=CKPT, failure_class="invalid_requirement", retryable=True,
            correction={"entries": [{"name": "contact_submission", "errors": DETAILS}]},
        ),
    )
    return run, repo, result


@pytest.mark.asyncio
async def test_a_runner_report_of_an_invalid_manifest_corrects_instead_of_ending_the_run(monkeypatch):
    run, repo, result = await runner_reports_blocked(monkeypatch, {"status": "correcting"})
    assert result == {"status": "correcting"}
    repo.update_status.assert_not_awaited()  # the run was NOT marked blocked
    # the failure was recorded before the decision, so the correction had something to act on
    assert run.plan["context"]["work_branch"]["failure_class"] == "invalid_requirement"


@pytest.mark.asyncio
async def test_when_correction_is_not_possible_the_run_ends_blocked_as_before(monkeypatch):
    _, repo, result = await runner_reports_blocked(monkeypatch, None)
    assert result == {"status": "blocked"}
    assert repo.update_status.await_args.args[1] == "blocked"