"""Tool for querying the Exploit Prediction Scoring System (EPSS)."""

from smolagents import Tool
import requests
import json

class EpsTool(Tool):
    """
    Tool for querying the Exploit Prediction Scoring System (EPSS).
    """
    
    name = "epss_search"
    description = """Tool for querying the Exploit Prediction Scoring System (EPSS).
    
    This tool allows obtaining EPSS scores for specific CVEs.
    Usage: epss_search(cve_id="CVE-2022-26332", date="2022-03-05")
    
    Example: epss_search(cve_id="CVE-2022-26332", date="2022-03-05")
    
    The EPSS score indicates the probability that a vulnerability will be exploited in the next 30 days.
    - Scores range from 0.0 to 1.0
    - Higher scores indicate higher probability of exploitation
    - Interpretation:
      * 0.7+ (70%+): High probability of exploitation
      * 0.4-0.7 (40-70%): Medium probability of exploitation  
      * 0.0-0.4 (0-40%): Low probability of exploitation"""
    
    inputs = {
        "cve_id": {
            "description": "The CVE ID to search for EPSS score (e.g., 'CVE-2022-26332').",
            "type": "string",
        },
        "date": {
            "description": "Optional date for the EPSS score (format: YYYY-MM-DD).",
            "type": "string",
            "nullable": True,
        },
    }
    
    output_type = "string"
    
    def __init__(self):
        super().__init__()
        self.base_url = "https://api.first.org/data/v1/epss"

    def forward(self, cve_id: str, date: str = None) -> str:
        """Search for EPSS score for a specific CVE."""
        try:
            # Build the URL with parameters
            params = {"cve": cve_id}
            if date:
                params["date"] = date
            
            url = f"{self.base_url}?{'&'.join([f'{k}={v}' for k, v in params.items()])}"
            
            response = requests.get(url)
            response.raise_for_status()
            data = response.json()
            
            if not data.get('data'):
                return f"No EPSS data found for {cve_id}"
            
            epss_data = data['data'][0]
            
            result = f"EPSS Score for {cve_id}:\n\n"
            result += f"- CVE ID: {epss_data.get('cve', 'Not available')}\n"
            result += f"- EPSS Score: {epss_data.get('epss', 'Not available')}\n"
            result += f"- Percentile: {epss_data.get('percentile', 'Not available')}\n"
            result += f"- Date: {epss_data.get('date', 'Not available')}\n"
            
            # Add interpretation
            epss_score = epss_data.get('epss')
            if epss_score is not None:
                try:
                    score = float(epss_score)
                    if score >= 0.7:
                        result += f"- Interpretation: High probability of exploitation ({score:.1%})\n"
                    elif score >= 0.4:
                        result += f"- Interpretation: Medium probability of exploitation ({score:.1%})\n"
                    else:
                        result += f"- Interpretation: Low probability of exploitation ({score:.1%})\n"
                except ValueError:
                    pass
            
            return result
            
        except requests.exceptions.RequestException as e:
            return f"Error accessing EPSS API: {str(e)}"
        except Exception as e:
            return f"Error processing EPSS data: {str(e)}"

    def _format_context(self, data: dict, cve_id: str) -> str:
        """Format the EPSS data for context."""
        if not data or not data.get('data'):
            return f"No EPSS data found for {cve_id}"
        
        epss_data = data['data'][0]
        context = []
        
        context.append(f"EPSS Score: {epss_data.get('epss', 'Not available')}")
        context.append(f"Percentile: {epss_data.get('percentile', 'Not available')}")
        context.append(f"Date: {epss_data.get('date', 'Not available')}")
        
        # Add interpretation
        epss_score = epss_data.get('epss')
        if epss_score is not None:
            try:
                score = float(epss_score)
                if score >= 0.7:
                    context.append(f"Risk Level: High ({score:.1%} probability of exploitation)")
                elif score >= 0.4:
                    context.append(f"Risk Level: Medium ({score:.1%} probability of exploitation)")
                else:
                    context.append(f"Risk Level: Low ({score:.1%} probability of exploitation)")
            except ValueError:
                pass
        
        return "\n".join(context)