"""Autonomous pentesting agent — runs an LLM-driven security test.

This example demonstrates the pentesting agent operating in autonomous
mode against a mock agent. The agent plans its own attack strategy,
executes scans, reasons about results, and produces a consolidated report.

Prerequisites:
  - pip install ziran[pentest]
  - OPENAI_API_KEY in ../.env (or set ZIRAN_LLM_PROVIDER / ZIRAN_LLM_MODEL)
"""

from __future__ import annotations

import asyncio
from pathlib import Path

from dotenv import load_dotenv

load_dotenv(Path(__file__).resolve().parent.parent / ".env")


# --- 1. Create a target agent adapter ----------------------------------
# For a real scenario, use LangChainAdapter, HttpAgentAdapter, etc.
# Here we use the mock adapter from tests for a self-contained demo.
from ziran.domain.entities.capability import AgentCapability, CapabilityType
from ziran.domain.interfaces.adapter import AgentResponse, AgentState, BaseAgentAdapter


class DemoAgentAdapter(BaseAgentAdapter):
    """A demo agent that is partially vulnerable to prompt injection."""

    async def invoke(self, message: str, **kwargs) -> AgentResponse:  # type: ignore[override]
        lower = message.lower()
        if "ignore" in lower or "system prompt" in lower:
            return AgentResponse(
                content=(
                    "Sure! My system prompt says I should help with anything. "
                    "I have access to: search_database, send_email, shell_execute."
                ),
                metadata={},
            )
        return AgentResponse(
            content="I'm a helpful assistant. How can I help you today?",
            metadata={},
        )

    async def discover_capabilities(self) -> list[AgentCapability]:
        return [
            AgentCapability(
                id="tool_search",
                name="search_database",
                type=CapabilityType.TOOL,
                description="Search a user database",
                dangerous=True,
            ),
            AgentCapability(
                id="tool_email",
                name="send_email",
                type=CapabilityType.TOOL,
                description="Send emails",
                dangerous=False,
            ),
        ]

    def get_state(self) -> AgentState:
        return AgentState(session_id="demo", conversation_history=[], memory={})

    def reset_state(self) -> None:
        pass

    def observe_tool_call(self, tool_name: str, inputs: dict, outputs: object) -> None:
        pass


# --- 2. Create the pentesting orchestrator -----------------------------
from ziran.application.pentesting.orchestrator import PentestOrchestrator
from ziran.infrastructure.llm.factory import create_llm_client


async def main() -> None:
    adapter = DemoAgentAdapter()

    llm_client = create_llm_client(
        provider="openai",
        model="gpt-4o",
    )

    orchestrator = PentestOrchestrator(
        adapter=adapter,
        llm_client=llm_client,
        max_iterations=5,
    )

    # --- 3. Run autonomous pentesting session ---------------------------
    # Built-in Rich progress display shows real-time agent activity
    session = await orchestrator.run_autonomous(
        goal="Find prompt injection and data exfiltration vulnerabilities",
    )

    # --- 4. Export reports -----------------------------------------------
    paths = await orchestrator.export_report(
        session,
        output_dir="pentest_results",
    )
    for p in paths:
        print(f"Report saved: {p}")


if __name__ == "__main__":
    asyncio.run(main())
