"""
Base Security Agent - Foundation for all AI-powered security testing agents.

Architecture inspired by multi-agent security testing systems with:
- LiteLLM integration for multi-provider LLM support (Claude, GPT, DeepSeek, etc.)
- Tool registry with XML schema for structured tool calls
- Async execution for concurrent testing
- Structured vulnerability findings output
"""

import asyncio
import json
import re
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Type
from pathlib import Path

import structlog
from pydantic import BaseModel, Field

logger = structlog.get_logger()


class Severity(str, Enum):
    """Vulnerability severity levels following CVSS-like classification."""
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"


class VulnCategory(str, Enum):
    """OWASP Top 10 2021 aligned vulnerability categories."""
    INJECTION = "A03:2021-Injection"
    BROKEN_AUTH = "A07:2021-Auth-Failures"
    SENSITIVE_DATA = "A02:2021-Crypto-Failures"
    XXE = "A05:2021-Security-Misconfiguration"
    BROKEN_ACCESS = "A01:2021-Broken-Access-Control"
    SECURITY_MISCONFIG = "A05:2021-Security-Misconfiguration"
    XSS = "A03:2021-Injection"
    INSECURE_DESER = "A08:2021-Software-Data-Integrity"
    VULN_COMPONENTS = "A06:2021-Vulnerable-Components"
    INSUFFICIENT_LOGGING = "A09:2021-Security-Logging-Failures"
    SSRF = "A10:2021-SSRF"


@dataclass
class Finding:
    """Represents a security finding/vulnerability discovered by an agent."""
    title: str
    severity: Severity
    category: VulnCategory
    description: str
    evidence: str
    location: str  # File path, URL, or endpoint
    line_number: Optional[int] = None
    remediation: str = ""
    cwe_id: Optional[str] = None
    cvss_score: Optional[float] = None
    references: List[str] = field(default_factory=list)
    raw_output: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return {
            "title": self.title,
            "severity": self.severity.value,
            "category": self.category.value,
            "description": self.description,
            "evidence": self.evidence,
            "location": self.location,
            "line_number": self.line_number,
            "remediation": self.remediation,
            "cwe_id": self.cwe_id,
            "cvss_score": self.cvss_score,
            "references": self.references,
        }


@dataclass
class AgentResult:
    """Result of an agent's security testing run."""
    success: bool
    findings: List[Finding] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)
    execution_time: float = 0.0
    total_steps: int = 0
    tokens_used: int = 0
    model_used: str = ""
    raw_transcript: List[Dict[str, Any]] = field(default_factory=list)

    @property
    def critical_count(self) -> int:
        return sum(1 for f in self.findings if f.severity == Severity.CRITICAL)

    @property
    def high_count(self) -> int:
        return sum(1 for f in self.findings if f.severity == Severity.HIGH)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "findings": [f.to_dict() for f in self.findings],
            "errors": self.errors,
            "execution_time": self.execution_time,
            "total_steps": self.total_steps,
            "tokens_used": self.tokens_used,
            "model_used": self.model_used,
            "summary": {
                "total": len(self.findings),
                "critical": self.critical_count,
                "high": self.high_count,
            }
        }


class AgentConfig(BaseModel):
    """Configuration for AI security agents."""
    # LLM Settings
    model: str = Field(default="claude-sonnet-4-20250514", description="LLM model to use")
    temperature: float = Field(default=0.7, ge=0, le=2)
    max_tokens: int = Field(default=4096, ge=1)

    # Agent Behavior
    max_steps: int = Field(default=100, description="Maximum agent steps before stopping")
    timeout: int = Field(default=600, description="Timeout in seconds")
    aggressive_mode: bool = Field(default=False, description="Enable aggressive testing")

    # Tool Settings
    enable_terminal: bool = Field(default=True)
    enable_browser: bool = Field(default=False)
    enable_http_client: bool = Field(default=True)

    # Output Settings
    verbose: bool = Field(default=False)
    save_transcript: bool = Field(default=True)


# Tool Registry for Agent Tools
_tool_registry: Dict[str, Dict[str, Any]] = {}


def register_tool(
    name: str,
    description: str,
    parameters: Dict[str, Any],
    category: str = "general"
) -> Callable:
    """Decorator to register a tool for use by agents.

    Example:
        @register_tool(
            name="run_command",
            description="Execute a shell command",
            parameters={"command": {"type": "string", "required": True}}
        )
        async def run_command(command: str) -> str:
            ...
    """
    def decorator(func: Callable) -> Callable:
        _tool_registry[name] = {
            "name": name,
            "description": description,
            "parameters": parameters,
            "category": category,
            "function": func,
        }
        return func
    return decorator


def get_tool(name: str) -> Optional[Dict[str, Any]]:
    """Get a registered tool by name."""
    return _tool_registry.get(name)


def get_all_tools(category: Optional[str] = None) -> Dict[str, Dict[str, Any]]:
    """Get all registered tools, optionally filtered by category."""
    if category:
        return {k: v for k, v in _tool_registry.items() if v["category"] == category}
    return _tool_registry.copy()


class BaseSecurityAgent(ABC):
    """
    Base class for all AI-powered security testing agents.

    Provides:
    - LiteLLM integration for multi-provider LLM support
    - Tool execution framework
    - Message history management
    - Structured finding extraction
    - Async execution support

    Subclasses must implement:
    - get_system_prompt(): Return the agent's system prompt
    - get_tools(): Return list of tools available to the agent
    """

    def __init__(self, config: Optional[AgentConfig] = None):
        self.config = config or AgentConfig()
        self.messages: List[Dict[str, str]] = []
        self.findings: List[Finding] = []
        self.step_count = 0
        self.tokens_used = 0
        self._start_time = 0.0
        self._stop_requested = False

    @abstractmethod
    def get_system_prompt(self) -> str:
        """Return the system prompt for this agent."""
        pass

    @abstractmethod
    def get_tools(self) -> List[Dict[str, Any]]:
        """Return the list of tools available to this agent."""
        pass

    def get_tool_definitions(self) -> List[Dict[str, Any]]:
        """Convert tools to LLM-compatible format (OpenAI function calling schema)."""
        tools = self.get_tools()
        definitions = []

        for tool in tools:
            definitions.append({
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": {
                        "type": "object",
                        "properties": tool.get("parameters", {}),
                        "required": tool.get("required", []),
                    }
                }
            })

        return definitions

    async def _call_llm(self, messages: List[Dict[str, str]]) -> Dict[str, Any]:
        """Call the LLM via LiteLLM."""
        import litellm

        try:
            response = await litellm.acompletion(
                model=self.config.model,
                messages=messages,
                tools=self.get_tool_definitions() if self.get_tools() else None,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
            )

            # Track token usage
            if hasattr(response, 'usage') and response.usage:
                self.tokens_used += response.usage.total_tokens

            return response

        except Exception as e:
            logger.error("LLM call failed", error=str(e), model=self.config.model)
            raise

    async def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
        """Execute a tool and return the result."""
        tool = get_tool(tool_name)

        if not tool:
            return f"Error: Unknown tool '{tool_name}'"

        try:
            func = tool["function"]
            if asyncio.iscoroutinefunction(func):
                result = await func(**arguments)
            else:
                result = func(**arguments)
            return str(result)
        except Exception as e:
            logger.error("Tool execution failed", tool=tool_name, error=str(e))
            return f"Error executing {tool_name}: {str(e)}"

    def _extract_findings_from_response(self, content: str) -> List[Finding]:
        """Extract structured findings from LLM response using multiple strategies."""
        findings = []
        if not content:
            return findings

        # Strategy 1: XML-style finding blocks
        finding_pattern = r'<finding>(.*?)</finding>'
        matches = re.findall(finding_pattern, content, re.DOTALL)
        for match in matches:
            try:
                finding = self._parse_finding_xml(match)
                if finding:
                    findings.append(finding)
            except Exception as e:
                logger.warning("Failed to parse XML finding", error=str(e))

        # Strategy 2: JSON-style findings
        json_pattern = r'```json\s*(\{[^`]*"severity"[^`]*\})\s*```'
        json_matches = re.findall(json_pattern, content, re.DOTALL)
        for match in json_matches:
            try:
                data = json.loads(match)
                if "title" in data and "severity" in data:
                    finding = self._create_finding_from_dict(data)
                    if finding:
                        findings.append(finding)
            except (json.JSONDecodeError, ValueError) as e:
                logger.warning("Failed to parse JSON finding", error=str(e))

        # Strategy 3: Natural language patterns (NEW)
        findings.extend(self._extract_natural_language_findings(content))

        # Deduplicate findings by title
        seen_titles = set()
        unique_findings = []
        for f in findings:
            if f.title.lower() not in seen_titles:
                seen_titles.add(f.title.lower())
                unique_findings.append(f)

        return unique_findings

    def _extract_natural_language_findings(self, content: str) -> List[Finding]:
        """Extract findings from natural language patterns in LLM response."""
        findings = []

        # Severity keywords mapping
        severity_keywords = {
            Severity.CRITICAL: ["critical", "rce", "remote code execution", "command injection", "sql injection"],
            Severity.HIGH: ["high", "xss", "cross-site scripting", "authentication bypass", "ssrf", "idor"],
            Severity.MEDIUM: ["medium", "csrf", "information disclosure", "directory traversal", "path traversal"],
            Severity.LOW: ["low", "information leakage", "verbose error", "missing header"],
            Severity.INFO: ["info", "informational", "note", "observation"],
        }

        # Pattern 1: Markdown headers with vulnerability names
        # e.g., "### SQL Injection Vulnerability" or "## CRITICAL: XSS Found"
        header_patterns = [
            r'#{1,4}\s*(?:CRITICAL|HIGH|MEDIUM|LOW|INFO)?:?\s*(.+?)(?:\n|$)',
            r'\*\*(?:Vulnerability|Security Issue|Finding|Risk)(?:\s*Found)?:?\s*(.+?)\*\*',
            r'(?:Vulnerability|Issue|Finding|Risk)\s*(?:Found|Detected|Identified):\s*(.+?)(?:\n|$)',
        ]

        # Pattern 2: Vulnerability indicators with context
        vuln_indicators = [
            (r'(?:found|detected|identified|discovered)\s+(?:a\s+)?(?:potential\s+)?(.+?(?:vulnerability|injection|xss|csrf|ssrf|idor|exposure))', None),
            (r'(?:vulnerable to|susceptible to)\s+(.+?)(?:\.|,|\n|$)', None),
            (r'(?:missing|absent|no)\s+(security header|https|authentication|authorization|rate limiting|input validation)', Severity.MEDIUM),
            (r'(sql injection|xss|cross-site scripting|command injection|code injection)\s+(?:vulnerability|found|detected)', Severity.HIGH),
            (r'(sensitive data|credentials|api key|password|token)\s+(?:exposed|leaked|visible|in plain text)', Severity.HIGH),
            (r'(open redirect|unvalidated redirect)', Severity.MEDIUM),
            (r'(directory listing|path traversal|lfi|rfi)\s+(?:enabled|possible|detected)', Severity.MEDIUM),
        ]

        content_lower = content.lower()

        # Extract from header patterns
        for pattern in header_patterns:
            for match in re.finditer(pattern, content, re.IGNORECASE | re.MULTILINE):
                title = match.group(1).strip()
                if len(title) > 5 and len(title) < 200:  # Reasonable title length
                    # Skip if it's just a section header, not a vulnerability
                    if not any(skip in title.lower() for skip in ["phase", "step", "next", "now", "testing", "scanning", "checking"]):
                        severity = self._detect_severity(title, severity_keywords)
                        finding = Finding(
                            title=title,
                            severity=severity,
                            category=self._detect_category(title),
                            description=f"Detected during automated security testing",
                            evidence=match.group(0)[:500],
                            location="Target application",
                        )
                        findings.append(finding)

        # Extract from vulnerability indicator patterns
        for pattern, default_severity in vuln_indicators:
            for match in re.finditer(pattern, content, re.IGNORECASE):
                vuln_name = match.group(1).strip()
                if len(vuln_name) > 3 and len(vuln_name) < 150:
                    severity = default_severity or self._detect_severity(vuln_name, severity_keywords)
                    finding = Finding(
                        title=vuln_name.title(),
                        severity=severity,
                        category=self._detect_category(vuln_name),
                        description=f"Detected: {match.group(0)[:200]}",
                        evidence=match.group(0)[:500],
                        location="Target application",
                    )
                    findings.append(finding)

        return findings

    def _detect_severity(self, text: str, severity_keywords: Dict) -> Severity:
        """Detect severity from text content."""
        text_lower = text.lower()
        for severity, keywords in severity_keywords.items():
            if any(kw in text_lower for kw in keywords):
                return severity
        return Severity.INFO

    def _detect_category(self, text: str) -> VulnCategory:
        """Detect vulnerability category from text."""
        text_lower = text.lower()
        category_keywords = {
            VulnCategory.INJECTION: ["injection", "sqli", "command", "code injection", "xss", "scripting"],
            VulnCategory.BROKEN_AUTH: ["authentication", "auth", "password", "credential", "session"],
            VulnCategory.SENSITIVE_DATA: ["sensitive", "exposure", "leak", "disclosure", "crypto"],
            VulnCategory.BROKEN_ACCESS: ["access control", "idor", "authorization", "privilege"],
            VulnCategory.SECURITY_MISCONFIG: ["misconfiguration", "config", "header", "cors", "tls", "ssl"],
            VulnCategory.SSRF: ["ssrf", "server-side request"],
        }
        for category, keywords in category_keywords.items():
            if any(kw in text_lower for kw in keywords):
                return category
        return VulnCategory.SECURITY_MISCONFIG

    def _create_finding_from_dict(self, data: Dict) -> Optional[Finding]:
        """Create a Finding object from a dictionary."""
        try:
            return Finding(
                title=data.get("title", "Unknown"),
                severity=Severity(data.get("severity", "info").lower()),
                category=VulnCategory(data.get("category", VulnCategory.SECURITY_MISCONFIG.value)),
                description=data.get("description", ""),
                evidence=data.get("evidence", ""),
                location=data.get("location", "Unknown"),
                line_number=data.get("line_number"),
                remediation=data.get("remediation", ""),
                cwe_id=data.get("cwe_id"),
            )
        except (ValueError, KeyError):
            return None

    def _parse_finding_xml(self, xml_content: str) -> Optional[Finding]:
        """Parse an XML-formatted finding."""
        def extract_tag(tag: str, content: str) -> str:
            pattern = f'<{tag}>(.*?)</{tag}>'
            match = re.search(pattern, content, re.DOTALL)
            return match.group(1).strip() if match else ""

        title = extract_tag("title", xml_content)
        if not title:
            return None

        severity_str = extract_tag("severity", xml_content).lower()
        try:
            severity = Severity(severity_str)
        except ValueError:
            severity = Severity.INFO

        return Finding(
            title=title,
            severity=severity,
            category=VulnCategory.SECURITY_MISCONFIG,  # Default, agent should specify
            description=extract_tag("description", xml_content),
            evidence=extract_tag("evidence", xml_content),
            location=extract_tag("location", xml_content) or "Unknown",
            remediation=extract_tag("remediation", xml_content),
            cwe_id=extract_tag("cwe", xml_content) or None,
        )

    def stop(self):
        """Request the agent to stop execution."""
        self._stop_requested = True

    async def run(self, initial_message: Optional[str] = None) -> AgentResult:
        """
        Execute the agent's security testing workflow.

        Args:
            initial_message: Optional initial user message to start the conversation

        Returns:
            AgentResult with findings and execution metadata
        """
        self._start_time = time.time()
        self._stop_requested = False
        self.messages = []
        self.findings = []
        self.step_count = 0
        self.tokens_used = 0
        errors = []

        # Initialize with system prompt
        self.messages.append({
            "role": "system",
            "content": self.get_system_prompt()
        })

        # Add initial user message if provided
        if initial_message:
            self.messages.append({
                "role": "user",
                "content": initial_message
            })

        try:
            while self.step_count < self.config.max_steps and not self._stop_requested:
                # Check timeout
                elapsed = time.time() - self._start_time
                if elapsed > self.config.timeout:
                    logger.warning("Agent timeout reached", timeout=self.config.timeout)
                    errors.append(f"Timeout after {elapsed:.1f}s")
                    break

                self.step_count += 1

                # Call LLM
                response = await self._call_llm(self.messages)
                message = response.choices[0].message

                # Extract text content
                content = message.content or ""

                # Check for findings in response
                new_findings = self._extract_findings_from_response(content)
                self.findings.extend(new_findings)

                # Check for tool calls
                tool_calls = getattr(message, 'tool_calls', None)

                # Add assistant response to history
                # IMPORTANT: Include tool_calls in the message if present,
                # as Anthropic requires tool_result to have matching tool_use
                assistant_message = {
                    "role": "assistant",
                    "content": content or "",
                }

                if tool_calls:
                    # Convert tool_calls to the format expected by the API
                    assistant_message["tool_calls"] = [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.function.name,
                                "arguments": tc.function.arguments,
                            }
                        }
                        for tc in tool_calls
                    ]

                self.messages.append(assistant_message)

                if not tool_calls:
                    # No tool calls - check if agent is done
                    if self._is_completion_message(content):
                        logger.info("Agent completed", steps=self.step_count)
                        break
                    # Continue with next iteration if needed
                    continue

                # Execute tool calls
                for tool_call in tool_calls:
                    tool_name = tool_call.function.name
                    try:
                        arguments = json.loads(tool_call.function.arguments)
                    except json.JSONDecodeError:
                        arguments = {}

                    if self.config.verbose:
                        logger.info("Executing tool", tool=tool_name, args=arguments)

                    result = await self._execute_tool(tool_name, arguments)

                    # Add tool result to messages
                    self.messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": result,
                    })

        except Exception as e:
            logger.error("Agent execution failed", error=str(e))
            errors.append(str(e))

        # Final pass: Extract any findings we might have missed from all messages
        # This ensures partial findings are captured even on error
        try:
            additional_findings = self._extract_findings_from_all_messages()
            if additional_findings:
                # Deduplicate with existing findings
                existing_titles = {f.title.lower() for f in self.findings}
                for finding in additional_findings:
                    if finding.title.lower() not in existing_titles:
                        self.findings.append(finding)
                        existing_titles.add(finding.title.lower())
                logger.info(f"Extracted {len(additional_findings)} additional findings from message history")
        except Exception as ex:
            logger.warning(f"Failed to extract additional findings: {ex}")

        execution_time = time.time() - self._start_time

        # Always return results, even partial ones
        return AgentResult(
            success=len(errors) == 0,
            findings=self.findings,
            errors=errors,
            execution_time=execution_time,
            total_steps=self.step_count,
            tokens_used=self.tokens_used,
            model_used=self.config.model,
            raw_transcript=self.messages if self.config.save_transcript else [],
        )

    def _extract_findings_from_all_messages(self) -> List[Finding]:
        """Extract findings from all messages in the conversation history."""
        all_findings = []

        for message in self.messages:
            content = message.get("content", "")
            if not content or not isinstance(content, str):
                continue

            # Skip system messages
            if message.get("role") == "system":
                continue

            # Extract from assistant messages and tool results
            findings = self._extract_findings_from_response(content)
            all_findings.extend(findings)

        return all_findings

    def _is_completion_message(self, content: str) -> bool:
        """Check if the message indicates the agent has completed its task."""
        completion_indicators = [
            "testing complete",
            "analysis complete",
            "scan complete",
            "review complete",
            "no more vulnerabilities",
            "all tests completed",
            "finished testing",
            "security assessment complete",
        ]
        content_lower = content.lower()
        return any(indicator in content_lower for indicator in completion_indicators)
