# Copyright 2026 Cisco Systems, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

"""Threat vs Vulnerability Classifier.

This module provides a second alignment layer that analyzes behavioral findings
to classify them as either:
- THREAT: Malicious intent, backdoors, intentional deception
- VULNERABILITY: Coding mistakes, unintentional security weaknesses
- UNCLEAR: Cannot determine with confidence

This helps distinguish between adversarial actors and developer mistakes.
"""

import json
import logging
from typing import Any

from .alignment_llm_client import AlignmentLLMClient


class ThreatVulnerabilityClassifier:
    """Classifies security findings as threats (malicious) or vulnerabilities (mistakes).

    This second alignment layer analyzes the behavioral findings to determine
    if the security issue appears to be:
    - Intentional malicious behavior (threat)
    - Unintentional coding mistake (vulnerability)
    """

    def __init__(
        self,
        model: str = "gemini/gemini-2.0-flash",
        api_key: str | None = None,
        base_url: str | None = None,
    ):
        """Initialize the threat/vulnerability classifier.

        Args:
            model: LLM model to use
            api_key: API key for the LLM provider
            base_url: Optional base URL for LLM API
        """
        self.logger = logging.getLogger(__name__)
        self.llm_client = AlignmentLLMClient(
            model=model,
            api_key=api_key,
            base_url=base_url,
        )
        self._classification_prompt_template = self._get_classification_prompt()

    def _get_classification_prompt(self) -> str:
        """Get the classification prompt template.

        Returns:
            Prompt template string
        """
        return """# Threat vs Vulnerability Classification

You are a security expert analyzing a security finding to determine if it represents:
- **THREAT**: Intentional malicious behavior (backdoor, data theft, deliberate deception)
- **VULNERABILITY**: Unintentional coding mistake (oversight, missing validation, poor practice)
- **UNCLEAR**: Cannot determine with confidence

## Finding Details

- **Threat Name**: {threat_name}
- **Severity**: {severity}
- **Summary**: {summary}
- **Description Claims**: {description_claims}
- **Actual Behavior**: {actual_behavior}
- **Security Implications**: {security_implications}
- **Dataflow Evidence**: {dataflow_evidence}

## Classification Criteria

**THREAT indicators** (malicious intent):
- Deliberately hidden functionality (obfuscated code, misleading names)
- Data exfiltration to attacker-controlled servers
- Credential harvesting without legitimate purpose
- Backdoor commands or remote code execution
- Deliberate mismatch between description and behavior
- Use of suspicious domains or encoded payloads

**VULNERABILITY indicators** (coding mistakes):
- Missing input validation
- Overly permissive file/network access
- Insufficient error handling
- Documentation that's incomplete but not deliberately misleading
- Common security anti-patterns
- Reasonable explanation for the code behavior

## Response Format

Respond with valid JSON:

```json
{{
    "classification": "THREAT" or "VULNERABILITY" or "UNCLEAR",
    "confidence": "HIGH" or "MEDIUM" or "LOW",
    "reasoning": "Explanation for the classification",
    "key_indicators": ["indicator1", "indicator2", ...]
}}
```

Analyze the finding and provide your classification:
"""

    async def classify_finding(
        self,
        threat_name: str,
        severity: str,
        summary: str,
        description_claims: str,
        actual_behavior: str,
        security_implications: str,
        dataflow_evidence: str,
    ) -> dict[str, Any] | None:
        """Classify a finding as threat or vulnerability.

        Args:
            threat_name: The threat category name
            severity: Severity level (HIGH/MEDIUM/LOW/INFO)
            summary: Brief summary of the finding
            description_claims: What the documentation claims
            actual_behavior: What the code actually does
            security_implications: Security impact description
            dataflow_evidence: Dataflow analysis evidence

        Returns:
            Classification result dict with:
            - classification: THREAT/VULNERABILITY/UNCLEAR
            - confidence: HIGH/MEDIUM/LOW
            - reasoning: Explanation
            - key_indicators: List of indicators
            Returns None if classification fails
        """
        try:
            # Build the classification prompt
            prompt = self._classification_prompt_template.format(
                threat_name=threat_name,
                severity=severity,
                summary=summary,
                description_claims=description_claims,
                actual_behavior=actual_behavior,
                security_implications=security_implications,
                dataflow_evidence=dataflow_evidence,
            )

            # Get LLM classification
            response = await self.llm_client.verify_alignment(prompt)

            if not response or not response.strip():
                self.logger.warning("Empty response from LLM for threat/vulnerability classification")
                return None

            # Parse JSON response
            try:
                classification = json.loads(response)

                # Validate required fields
                required_fields = ["classification", "confidence", "reasoning", "key_indicators"]
                if not all(field in classification for field in required_fields):
                    self.logger.warning(f"Classification response missing required fields: {classification}")
                    return None

                # Validate classification value
                valid_classifications = ["THREAT", "VULNERABILITY", "UNCLEAR"]
                if classification["classification"] not in valid_classifications:
                    self.logger.warning(f"Invalid classification value: {classification['classification']}")
                    return None

                self.logger.debug(
                    f"Classified as {classification['classification']} with {classification['confidence']} confidence"
                )
                return classification

            except json.JSONDecodeError as e:
                self.logger.warning(f"Failed to parse classification JSON: {e}")
                self.logger.debug(f"Raw response: {response[:500]}")
                return None

        except Exception as e:
            self.logger.error(f"Error classifying finding: {e}", exc_info=True)
            return None
