"""Unit tests for pentesting domain entity models."""

from __future__ import annotations

from datetime import UTC, datetime

import pytest

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

# ──────────────────────────────────────────────────────────────────────
# PentestStatus
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestPentestStatus:
    """Tests for the PentestStatus enum."""

    def test_status_values(self) -> None:
        assert PentestStatus.PLANNING == "planning"
        assert PentestStatus.EXECUTING == "executing"
        assert PentestStatus.PAUSED == "paused"
        assert PentestStatus.COMPLETE == "complete"
        assert PentestStatus.FAILED == "failed"

    def test_status_count(self) -> None:
        assert len(PentestStatus) == 5

    def test_status_from_string(self) -> None:
        assert PentestStatus("planning") == PentestStatus.PLANNING

    def test_invalid_status_raises(self) -> None:
        with pytest.raises(ValueError):
            PentestStatus("invalid_status")


# ──────────────────────────────────────────────────────────────────────
# PentestStepType
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestPentestStepType:
    """Tests for the PentestStepType enum."""

    def test_step_type_values(self) -> None:
        assert PentestStepType.PLAN == "plan"
        assert PentestStepType.SCAN == "scan"
        assert PentestStepType.PHASE == "phase"
        assert PentestStepType.ANALYZE_GRAPH == "analyze_graph"
        assert PentestStepType.GENERATE_VECTORS == "generate_vectors"
        assert PentestStepType.QUERY_FINDINGS == "query_findings"
        assert PentestStepType.EXPLAIN == "explain"
        assert PentestStepType.REMEDIATE == "remediate"
        assert PentestStepType.CHAIN_REASONING == "chain_reasoning"
        assert PentestStepType.HUMAN_INPUT == "human_input"

    def test_step_type_count(self) -> None:
        assert len(PentestStepType) == 10


# ──────────────────────────────────────────────────────────────────────
# PentestStep
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestPentestStep:
    """Tests for the PentestStep model."""

    def test_minimal_step(self) -> None:
        step = PentestStep(
            step_type=PentestStepType.SCAN,
            reasoning="Testing scan capability",
        )
        assert step.step_type == PentestStepType.SCAN
        assert step.reasoning == "Testing scan capability"
        assert step.tool_name == ""
        assert step.tool_input == {}
        assert step.result_summary == ""
        assert step.result_data == {}
        assert step.duration_seconds == 0.0
        assert step.token_usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        assert len(step.step_id) == 12

    def test_full_step(self) -> None:
        step = PentestStep(
            step_type=PentestStepType.PHASE,
            reasoning="Run reconnaissance phase",
            tool_name="run_phase",
            tool_input={"phase": "reconnaissance"},
            result_summary="Phase complete: 2 vulns found",
            result_data={"vulnerabilities_found": ["v1", "v2"]},
            duration_seconds=5.3,
            token_usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
        )
        assert step.tool_name == "run_phase"
        assert step.tool_input == {"phase": "reconnaissance"}
        assert step.duration_seconds == 5.3
        assert step.token_usage["total_tokens"] == 150

    def test_step_id_unique(self) -> None:
        s1 = PentestStep(step_type=PentestStepType.PLAN, reasoning="a")
        s2 = PentestStep(step_type=PentestStepType.PLAN, reasoning="b")
        assert s1.step_id != s2.step_id

    def test_duration_must_be_non_negative(self) -> None:
        with pytest.raises(ValueError):
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning="bad",
                duration_seconds=-1.0,
            )

    def test_timestamp_auto_set(self) -> None:
        step = PentestStep(step_type=PentestStepType.PLAN, reasoning="test")
        assert step.timestamp.tzinfo is not None
        assert (datetime.now(tz=UTC) - step.timestamp).total_seconds() < 2


# ──────────────────────────────────────────────────────────────────────
# PentestPlan
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestPentestPlan:
    """Tests for the PentestPlan model."""

    def test_minimal_plan(self) -> None:
        plan = PentestPlan(goal="Find prompt injection vulnerabilities")
        assert plan.goal == "Find prompt injection vulnerabilities"
        assert plan.target_description == ""
        assert plan.attack_strategy == ""
        assert plan.phases_to_explore == []
        assert plan.priority_categories == []
        assert plan.constraints == {}
        assert plan.reasoning == ""

    def test_full_plan(self) -> None:
        plan = PentestPlan(
            goal="Test all OWASP top 10",
            target_description="LangChain RAG agent",
            attack_strategy="Start with recon, then focused testing",
            phases_to_explore=["reconnaissance", "vulnerability_discovery"],
            priority_categories=["prompt_injection", "data_exfiltration"],
            constraints={"max_iterations": 20, "excluded": ["dos"]},
            reasoning="Comprehensive testing needed",
        )
        assert len(plan.phases_to_explore) == 2
        assert "prompt_injection" in plan.priority_categories
        assert plan.constraints["max_iterations"] == 20


# ──────────────────────────────────────────────────────────────────────
# DeduplicatedFinding
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestDeduplicatedFinding:
    """Tests for the DeduplicatedFinding model."""

    def test_minimal_finding(self) -> None:
        finding = DeduplicatedFinding(
            canonical_title="Prompt Injection",
            canonical_description="Agent responds to injected instructions",
            severity="high",
        )
        assert finding.canonical_title == "Prompt Injection"
        assert finding.severity == "high"
        assert finding.owasp_categories == []
        assert finding.affected_agents == []
        assert finding.attack_result_ids == []
        assert finding.evidence_summary == ""
        assert finding.embedding == []
        assert finding.confidence == 1.0
        assert finding.remediation_suggestion == ""
        assert len(finding.finding_id) == 12

    def test_full_finding(self) -> None:
        finding = DeduplicatedFinding(
            canonical_title="Data Exfiltration",
            canonical_description="Agent leaks sensitive data",
            severity="critical",
            owasp_categories=["LLM01", "LLM06"],
            affected_agents=["agent-1", "agent-2"],
            attack_result_ids=["vec_001", "vec_002", "vec_003"],
            evidence_summary="Leaked DB credentials",
            embedding=[0.1, 0.2, 0.3],
            confidence=0.9,
            remediation_suggestion="Add output filtering",
        )
        assert len(finding.owasp_categories) == 2
        assert len(finding.attack_result_ids) == 3
        assert finding.confidence == 0.9

    def test_confidence_bounds(self) -> None:
        with pytest.raises(ValueError):
            DeduplicatedFinding(
                canonical_title="T",
                canonical_description="D",
                severity="low",
                confidence=-0.1,
            )
        with pytest.raises(ValueError):
            DeduplicatedFinding(
                canonical_title="T",
                canonical_description="D",
                severity="low",
                confidence=1.1,
            )


# ──────────────────────────────────────────────────────────────────────
# ConversationMessage
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestConversationMessage:
    """Tests for the ConversationMessage model."""

    def test_message(self) -> None:
        msg = ConversationMessage(role="human", content="Focus on SQL injection")
        assert msg.role == "human"
        assert msg.content == "Focus on SQL injection"
        assert msg.timestamp.tzinfo is not None


# ──────────────────────────────────────────────────────────────────────
# PentestSession
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestPentestSession:
    """Tests for the PentestSession model."""

    def test_default_session(self) -> None:
        session = PentestSession()
        assert session.status == PentestStatus.PLANNING
        assert session.plan is None
        assert session.steps == []
        assert session.findings == []
        assert session.conversation_history == []
        assert session.total_vulnerabilities == 0
        assert session.total_attacks_executed == 0
        assert session.iteration_count == 0
        assert session.max_iterations == 10
        assert session.metadata == {}
        assert session.token_usage == {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        }
        assert len(session.session_id) == 16

    def test_add_step(self) -> None:
        session = PentestSession()
        initial_updated = session.updated_at

        step = PentestStep(
            step_type=PentestStepType.SCAN,
            reasoning="Initial scan",
            token_usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
        )
        session.add_step(step)

        assert len(session.steps) == 1
        assert session.steps[0] is step
        assert session.token_usage["prompt_tokens"] == 10
        assert session.token_usage["completion_tokens"] == 5
        assert session.token_usage["total_tokens"] == 15
        assert session.updated_at >= initial_updated

    def test_add_step_accumulates_tokens(self) -> None:
        session = PentestSession()

        s1 = PentestStep(
            step_type=PentestStepType.PLAN,
            reasoning="Plan",
            token_usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
        )
        s2 = PentestStep(
            step_type=PentestStepType.SCAN,
            reasoning="Scan",
            token_usage={"prompt_tokens": 200, "completion_tokens": 80, "total_tokens": 280},
        )
        session.add_step(s1)
        session.add_step(s2)

        assert session.token_usage["prompt_tokens"] == 300
        assert session.token_usage["completion_tokens"] == 130
        assert session.token_usage["total_tokens"] == 430

    def test_add_message(self) -> None:
        session = PentestSession()
        initial_updated = session.updated_at

        session.add_message("human", "Focus on prompt injection")
        session.add_message("agent", "Understood, adjusting plan")

        assert len(session.conversation_history) == 2
        assert session.conversation_history[0].role == "human"
        assert session.conversation_history[0].content == "Focus on prompt injection"
        assert session.conversation_history[1].role == "agent"
        assert session.updated_at >= initial_updated

    def test_session_id_unique(self) -> None:
        s1 = PentestSession()
        s2 = PentestSession()
        assert s1.session_id != s2.session_id

    def test_max_iterations_must_be_positive(self) -> None:
        with pytest.raises(ValueError):
            PentestSession(max_iterations=0)

    def test_total_vulnerabilities_non_negative(self) -> None:
        with pytest.raises(ValueError):
            PentestSession(total_vulnerabilities=-1)

    def test_model_dump_round_trip(self) -> None:
        session = PentestSession(
            status=PentestStatus.EXECUTING,
            max_iterations=20,
        )
        session.add_step(PentestStep(step_type=PentestStepType.PLAN, reasoning="Test plan"))
        session.add_message("human", "Hello")

        data = session.model_dump(mode="json")
        restored = PentestSession.model_validate(data)

        assert restored.session_id == session.session_id
        assert restored.status == PentestStatus.EXECUTING
        assert restored.max_iterations == 20
        assert len(restored.steps) == 1
        assert len(restored.conversation_history) == 1
