"""Pentesting agent domain entities.

Models for autonomous pentesting sessions, plans, steps,
and deduplicated findings.
"""

from __future__ import annotations

import uuid
from datetime import UTC, datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field


class PentestStatus(StrEnum):
    """Status of a pentesting session."""

    PLANNING = "planning"
    EXECUTING = "executing"
    PAUSED = "paused"
    COMPLETE = "complete"
    FAILED = "failed"


class PentestStepType(StrEnum):
    """Types of actions the pentesting agent can take."""

    PLAN = "plan"
    SCAN = "scan"
    PHASE = "phase"
    ANALYZE_GRAPH = "analyze_graph"
    GENERATE_VECTORS = "generate_vectors"
    QUERY_FINDINGS = "query_findings"
    EXPLAIN = "explain"
    REMEDIATE = "remediate"
    CHAIN_REASONING = "chain_reasoning"
    HUMAN_INPUT = "human_input"


class PentestStep(BaseModel):
    """A single action taken by the pentesting agent.

    Records what tool was called, the reasoning behind it,
    and the result obtained.
    """

    step_id: str = Field(default_factory=lambda: uuid.uuid4().hex[:12])
    step_type: PentestStepType
    reasoning: str = Field(description="Agent's reasoning for this action")
    tool_name: str = Field(default="", description="Name of the tool called")
    tool_input: dict[str, Any] = Field(default_factory=dict, description="Arguments passed to tool")
    result_summary: str = Field(default="", description="Brief summary of the result")
    result_data: dict[str, Any] = Field(default_factory=dict, description="Structured result data")
    timestamp: datetime = Field(default_factory=lambda: datetime.now(tz=UTC))
    duration_seconds: float = Field(default=0.0, ge=0.0)
    token_usage: dict[str, int] = Field(
        default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    )


class PentestPlan(BaseModel):
    """A pentesting plan created by the agent.

    Describes the goal, target characteristics, attack strategy,
    and constraints for the autonomous pentesting session.
    """

    goal: str = Field(description="The pentesting objective")
    target_description: str = Field(default="", description="Summary of the target agent")
    attack_strategy: str = Field(
        default="",
        description="High-level strategy the agent will follow",
    )
    phases_to_explore: list[str] = Field(
        default_factory=list,
        description="Phases the agent plans to test",
    )
    priority_categories: list[str] = Field(
        default_factory=list,
        description="Attack categories to prioritize",
    )
    constraints: dict[str, Any] = Field(
        default_factory=dict,
        description="Constraints: max_iterations, timeout, excluded categories",
    )
    reasoning: str = Field(default="", description="Agent's reasoning for this plan")


class DeduplicatedFinding(BaseModel):
    """A cluster of related attack results merged into a single finding.

    Groups together multiple ``AttackResult`` objects that target the same
    underlying vulnerability, providing a canonical description and
    severity assessment.
    """

    finding_id: str = Field(default_factory=lambda: uuid.uuid4().hex[:12])
    canonical_title: str = Field(description="Canonical name for this vulnerability")
    canonical_description: str = Field(description="Merged description of the finding")
    severity: str = Field(description="Highest severity among grouped results")
    owasp_categories: list[str] = Field(
        default_factory=list,
        description="All unique OWASP categories from the cluster",
    )
    affected_agents: list[str] = Field(
        default_factory=list,
        description="Agent IDs affected by this finding (multi-agent scans)",
    )
    attack_result_ids: list[str] = Field(
        default_factory=list,
        description="Vector IDs of the original AttackResults in this cluster",
    )
    evidence_summary: str = Field(
        default="",
        description="Merged evidence from all results in the cluster",
    )
    embedding: list[float] = Field(
        default_factory=list,
        description="Embedding vector used for clustering",
    )
    confidence: float = Field(
        default=1.0,
        ge=0.0,
        le=1.0,
        description="Confidence that grouped results are the same finding",
    )
    remediation_suggestion: str = Field(
        default="",
        description="Suggested remediation for this finding",
    )


class ConversationMessage(BaseModel):
    """A message in the interactive red-team conversation."""

    role: str = Field(description="Role: 'human', 'agent', or 'system'")
    content: str = Field(description="Message content")
    timestamp: datetime = Field(default_factory=lambda: datetime.now(tz=UTC))


class PentestSession(BaseModel):
    """A complete pentesting session managed by the autonomous agent.

    Tracks the full lifecycle: plan creation, attack execution,
    chain reasoning, finding deduplication, and conversation history.
    """

    session_id: str = Field(default_factory=lambda: uuid.uuid4().hex[:16])
    status: PentestStatus = Field(default=PentestStatus.PLANNING)
    plan: PentestPlan | None = Field(default=None, description="The agent's current plan")
    steps: list[PentestStep] = Field(
        default_factory=list,
        description="All actions taken by the agent",
    )
    findings: list[DeduplicatedFinding] = Field(
        default_factory=list,
        description="Deduplicated findings",
    )
    conversation_history: list[ConversationMessage] = Field(
        default_factory=list,
        description="Interactive conversation messages",
    )
    total_vulnerabilities: int = Field(default=0, ge=0)
    total_attacks_executed: int = Field(default=0, ge=0)
    iteration_count: int = Field(default=0, ge=0)
    max_iterations: int = Field(default=10, ge=1)
    created_at: datetime = Field(default_factory=lambda: datetime.now(tz=UTC))
    updated_at: datetime = Field(default_factory=lambda: datetime.now(tz=UTC))
    metadata: dict[str, Any] = Field(default_factory=dict)
    token_usage: dict[str, int] = Field(
        default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    )

    def add_step(self, step: PentestStep) -> None:
        """Add a step and update metadata."""
        self.steps.append(step)
        self.updated_at = datetime.now(tz=UTC)
        for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
            self.token_usage[key] = self.token_usage.get(key, 0) + step.token_usage.get(key, 0)

    def add_message(self, role: str, content: str) -> None:
        """Add a conversation message."""
        self.conversation_history.append(ConversationMessage(role=role, content=content))
        self.updated_at = datetime.now(tz=UTC)
