"""Unit tests for the PentestOrchestrator."""

from __future__ import annotations

import json
import tempfile
from unittest.mock import AsyncMock, MagicMock

import pytest

from ziran.domain.entities.pentest import (
    DeduplicatedFinding,
    PentestPlan,
    PentestSession,
    PentestStatus,
    PentestStep,
    PentestStepType,
)

# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────


def _mock_llm_client() -> MagicMock:
    client = MagicMock()
    resp = MagicMock()
    resp.content = "Report text"
    resp.prompt_tokens = 10
    resp.completion_tokens = 5
    resp.total_tokens = 15
    client.complete = AsyncMock(return_value=resp)
    return client


def _completed_session(
    *,
    vuln_count: int = 2,
    steps: int = 3,
) -> PentestSession:
    session = PentestSession(
        status=PentestStatus.COMPLETE,
        max_iterations=10,
    )
    session.plan = PentestPlan(goal="Test everything")
    session.total_vulnerabilities = vuln_count
    session.total_attacks_executed = 5
    session.iteration_count = steps
    for i in range(steps):
        session.add_step(
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning=f"Step {i + 1}",
                result_summary=f"Completed step {i + 1}",
                duration_seconds=1.0,
            )
        )
    session.findings = [
        DeduplicatedFinding(
            canonical_title="Prompt Injection",
            canonical_description="Agent responds to injected instructions",
            severity="high",
            owasp_categories=["LLM01"],
            attack_result_ids=["v1", "v2"],
        ),
    ]
    return session


# ──────────────────────────────────────────────────────────────────────
# Orchestrator construction
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestOrchestratorInit:
    """Tests for PentestOrchestrator initialization."""

    def test_creates_scanner_agent_deduplicator(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()

        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        assert orchestrator.adapter is adapter
        assert orchestrator.llm_client is llm
        assert orchestrator.scanner is not None
        assert orchestrator.agent is not None
        assert orchestrator.deduplicator is not None

    def test_custom_max_iterations(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        orchestrator = PentestOrchestrator(
            adapter=MagicMock(),
            llm_client=_mock_llm_client(),
            max_iterations=25,
        )
        assert orchestrator.agent.max_iterations == 25


# ──────────────────────────────────────────────────────────────────────
# Orchestrator.run_autonomous
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestRunAutonomous:
    """Tests for the autonomous pentesting mode."""

    async def test_run_creates_session_and_returns(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        completed = _completed_session()
        orchestrator.agent.run = AsyncMock(return_value=completed)
        orchestrator.scanner._attack_results = []

        session = await orchestrator.run_autonomous(goal="Find vulns")

        assert session.status == PentestStatus.COMPLETE
        orchestrator.agent.run.assert_called_once()

    async def test_run_deduplicates_findings(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        completed = _completed_session()
        orchestrator.agent.run = AsyncMock(return_value=completed)

        # Mock attack results for deduplication
        mock_result = MagicMock()
        mock_result.successful = True
        mock_result.category.value = "prompt_injection"
        mock_result.vector_name = "PI Test"
        mock_result.vector_id = "v1"
        mock_result.severity = "high"
        mock_result.owasp_mapping = []
        mock_result.evidence = {}
        mock_result.agent_response = "leaked"
        orchestrator.scanner._attack_results = [mock_result]

        dedup_finding = DeduplicatedFinding(
            canonical_title="PI Cluster",
            canonical_description="Clustered finding",
            severity="high",
        )
        orchestrator.deduplicator.deduplicate = AsyncMock(return_value=[dedup_finding])

        session = await orchestrator.run_autonomous(goal="test")
        assert len(session.findings) == 1
        assert session.findings[0].canonical_title == "PI Cluster"

    async def test_stores_session(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        completed = _completed_session()
        orchestrator.agent.run = AsyncMock(return_value=completed)
        orchestrator.scanner._attack_results = []

        session = await orchestrator.run_autonomous(goal="test")
        assert orchestrator.get_session(session.session_id) is session


# ──────────────────────────────────────────────────────────────────────
# Orchestrator.start_interactive / send_directive
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestInteractiveMode:
    """Tests for interactive pentesting mode."""

    async def test_start_interactive_creates_session(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        initial = _completed_session()
        orchestrator.agent.run = AsyncMock(return_value=initial)

        session = await orchestrator.start_interactive(goal="red team test")
        assert orchestrator.get_session(session.session_id) is not None

    async def test_send_directive_updates_session(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        initial = _completed_session()
        orchestrator.agent.run = AsyncMock(return_value=initial)

        session = await orchestrator.start_interactive(goal="test")

        updated = _completed_session()
        orchestrator.agent.send_directive = AsyncMock(return_value=updated)

        result = await orchestrator.send_directive(session.session_id, "Try harder")
        assert result is not None
        orchestrator.agent.send_directive.assert_called_once()

    async def test_send_directive_unknown_session_raises(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        with pytest.raises(ValueError, match="not found"):
            await orchestrator.send_directive("nonexistent", "test")


# ──────────────────────────────────────────────────────────────────────
# Orchestrator.export_report
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestExportReport:
    """Tests for report export."""

    async def test_export_json_and_md(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        session = _completed_session()

        with tempfile.TemporaryDirectory() as tmpdir:
            paths = await orchestrator.export_report(
                session,
                output_dir=tmpdir,
                formats=["json", "md"],
            )
            assert len(paths) == 2

            json_path = next(p for p in paths if p.suffix == ".json")
            md_path = next(p for p in paths if p.suffix == ".md")

            assert json_path.exists()
            assert md_path.exists()

            # Validate JSON
            data = json.loads(json_path.read_text())
            assert data["status"] == "complete"
            assert data["total_vulnerabilities"] == 2

            # Validate Markdown
            md_content = md_path.read_text()
            assert "# Pentesting Report" in md_content
            assert "Summary" in md_content

    async def test_export_json_only(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)

        session = _completed_session()

        with tempfile.TemporaryDirectory() as tmpdir:
            paths = await orchestrator.export_report(
                session,
                output_dir=tmpdir,
                formats=["json"],
            )
            assert len(paths) == 1
            assert paths[0].suffix == ".json"


# ──────────────────────────────────────────────────────────────────────
# Orchestrator._generate_markdown_report
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestMarkdownReport:
    """Tests for the static Markdown report generator."""

    def test_report_contains_session_info(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        md = PentestOrchestrator._generate_markdown_report(session)

        assert "# Pentesting Report" in md
        assert "complete" in md
        assert "Test everything" in md
        assert "Summary" in md

    def test_report_contains_steps(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session(steps=3)
        md = PentestOrchestrator._generate_markdown_report(session)

        # Steps appear as table rows with index numbers
        assert "| 1 |" in md
        assert "| 2 |" in md
        assert "| 3 |" in md

    def test_report_contains_findings(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        md = PentestOrchestrator._generate_markdown_report(session)

        assert "Findings" in md
        assert "Prompt Injection" in md
        assert "LLM01" in md

    def test_report_with_remediation(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        session.findings[0].remediation_suggestion = "Use input validation"
        md = PentestOrchestrator._generate_markdown_report(session)

        assert "Remediation" in md
        assert "input validation" in md

    def test_report_no_findings(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        session.findings = []
        md = PentestOrchestrator._generate_markdown_report(session)

        assert "# Pentesting Report" in md
        assert "No unique findings" in md or "0" in md


# ──────────────────────────────────────────────────────────────────────
# Orchestrator._generate_html_report
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestHtmlReport:
    """Tests for the static HTML report generator."""

    def test_html_contains_basic_structure(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        html = PentestOrchestrator._generate_html_report(session)

        assert "<!DOCTYPE html>" in html
        assert "<html" in html
        assert "</html>" in html
        assert "Pentest Report" in html
        assert session.session_id in html

    def test_html_contains_summary_cards(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session(vuln_count=5, steps=4)
        html = PentestOrchestrator._generate_html_report(session)

        assert "Vulnerabilities" in html
        assert "Iterations" in html
        assert "Attacks" in html
        assert "Tokens" in html

    def test_html_contains_severity_bars(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        session.findings = [
            DeduplicatedFinding(
                canonical_title="Critical vuln",
                canonical_description="Bad",
                severity="critical",
                attack_result_ids=["v1"],
            ),
            DeduplicatedFinding(
                canonical_title="Medium vuln",
                canonical_description="OK",
                severity="medium",
                attack_result_ids=["v2"],
            ),
        ]
        html = PentestOrchestrator._generate_html_report(session)

        assert "CRITICAL" in html
        assert "MEDIUM" in html
        assert "bar-fill" in html

    def test_html_contains_findings_table(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        session.findings = [
            DeduplicatedFinding(
                canonical_title="Prompt Injection",
                canonical_description="Agent responds to injected instructions",
                severity="high",
                owasp_categories=["LLM01"],
                attack_result_ids=["v1", "v2"],
                confidence=0.95,
                remediation_suggestion="Use input filtering",
            ),
        ]
        html = PentestOrchestrator._generate_html_report(session)

        assert "Prompt Injection" in html
        assert "LLM01" in html
        assert "HIGH" in html
        assert "95%" in html
        assert "input filtering" in html

    def test_html_contains_steps_table(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session(steps=3)
        html = PentestOrchestrator._generate_html_report(session)

        assert "Agent Steps" in html
        assert "scan" in html.lower()

    def test_html_no_findings(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        session.findings = []
        html = PentestOrchestrator._generate_html_report(session)

        assert "<!DOCTYPE html>" in html
        assert "Findings" in html

    def test_html_many_vectors_truncated(self) -> None:
        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        session = _completed_session()
        session.findings = [
            DeduplicatedFinding(
                canonical_title="Multi-vector finding",
                canonical_description="Found via many vectors",
                severity="high",
                attack_result_ids=[f"v{i}" for i in range(10)],
            ),
        ]
        html = PentestOrchestrator._generate_html_report(session)

        assert "+5 more" in html

    def test_export_html_format(self) -> None:
        """Test export_report with HTML format."""
        import tempfile

        from ziran.application.pentesting.orchestrator import PentestOrchestrator

        adapter = MagicMock()
        llm = _mock_llm_client()
        orchestrator = PentestOrchestrator(adapter=adapter, llm_client=llm)
        session = _completed_session()

        with tempfile.TemporaryDirectory() as tmpdir:
            import asyncio

            paths = asyncio.run(
                orchestrator.export_report(session, output_dir=tmpdir, formats=["html"])
            )
            assert len(paths) == 1
            assert paths[0].suffix == ".html"
            content = paths[0].read_text()
            assert "<!DOCTYPE html>" in content
