import aiohttp
import asyncio
import json
import os
from datetime import datetime, timedelta

class CVEDataFetcher:
    """Module to fetch CVE data from NVD or other CVE databases."""
    
    def __init__(self, cache_dir='.cve_cache'):
        self.cache_dir = cache_dir
        self.cache_file = os.path.join(cache_dir, 'cve_data.json')
        self.cache_expiry = timedelta(days=1)
        self.nvd_api_url = "https://services.nvd.nist.gov/rest/json/cves/1.0"
        
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
    
    async def fetch_cve_data(self, start_index=0, results_per_page=2000):
        """Fetch CVE data from NVD API asynchronously."""
        if self._is_cache_valid():
            return self._load_cache()
        
        all_cve_items = []
        
        async with aiohttp.ClientSession() as session:
            while True:
                params = {
                    'startIndex': start_index,
                    'resultsPerPage': results_per_page
                }
                
                try:
                    async with session.get(self.nvd_api_url, params=params) as response:
                        response.raise_for_status()
                        data = await response.json()
                
                    cve_items = data.get('result', {}).get('CVE_Items', [])
                    all_cve_items.extend(cve_items)
                
                    if len(cve_items) < results_per_page:
                        break
                
                    start_index += results_per_page
                except aiohttp.ClientError as e:
                    print(f"Error fetching CVE data: {e}")
                    break
        
        self._cache_data(all_cve_items)
        return all_cve_items
    
    def _is_cache_valid(self):
        """Check if the cache is still valid."""
        if not os.path.exists(self.cache_file):
            return False
        
        file_time = datetime.fromtimestamp(os.path.getmtime(self.cache_file))
        return datetime.now() - file_time < self.cache_expiry
    
    def _load_cache(self):
        """Load CVE data from cache."""
        with open(self.cache_file, 'r') as f:
            return json.load(f)
    
    def _cache_data(self, data):
        """Cache the fetched CVE data."""
        with open(self.cache_file, 'w') as f:
            json.dump(data, f)
