use std::path::{Path, PathBuf};
use std::fs::{create_dir_all, File};
use std::io::{self, Read, Write};
use reqwest::Client;
use flate2::read::GzDecoder;
use serde_json::Value;
use clap::{Parser, ValueEnum};
use chrono::{Utc, Datelike};
use crate::error::CollectorError;

#[derive(Copy, Clone, PartialEq, Eq, ValueEnum, Debug)]
pub enum FeedType {
    Recent,
    Full,
}

pub struct NvdCveFetcher {
    client: Client,
    download_dir: PathBuf,
}

impl NvdCveFetcher {
    pub fn new() -> io::Result<Self> {
        let download_dir = Path::new("nvd_feeds");
        create_dir_all(&download_dir)?;

        Ok(Self {
            client: Client::new(),
            download_dir: download_dir.to_path_buf(),
        })
    }

    pub async fn fetch_cves(&self, feed_type: FeedType) -> Result<Vec<String>, CollectorError> {
        match feed_type {
            FeedType::Recent => self.fetch_recent_cves().await,
            FeedType::Full => self.fetch_full_historical_cves().await,
        }
    }

    async fn fetch_recent_cves(&self) -> Result<Vec<String>, CollectorError> {
        let modified_url = "https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-modified.json.gz";
        let recent_url = "https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-recent.json.gz";

        let mut all_cve_ids = Vec::new();

        // Download and process modified feed
        let modified_filepath = self.download_feed(modified_url).await?;
        let modified_data = self.decompress_feed(&modified_filepath)?;
        let modified_cves = self.extract_cve_ids(&modified_data)?;
        all_cve_ids.extend(modified_cves);

        // Download and process recent feed
        let recent_filepath = self.download_feed(recent_url).await?;
        let recent_data = self.decompress_feed(&recent_filepath)?;
        let recent_cves = self.extract_cve_ids(&recent_data)?;
        all_cve_ids.extend(recent_cves);

        // Remove duplicates
        all_cve_ids.sort_unstable();
        all_cve_ids.dedup();

        Ok(all_cve_ids)
    }

    async fn fetch_full_historical_cves(&self) -> Result<Vec<String>, CollectorError> {
        let current_year = Utc::now().year();
        let mut all_cve_ids = Vec::new();

        for year in 2002..=current_year {
            let url = format!(
                "https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-{}.json.gz",
                year
            );

            println!("Downloading CVE feed for year: {}", year);

            // Download and process the feed
            let filepath = self.download_feed(&url).await?;
            let data = self.decompress_feed(&filepath)?;
            let cves = self.extract_cve_ids(&data)?;
            
            all_cve_ids.extend(cves);

            // Optional: Add a small delay to prevent overwhelming the server
            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
        }

        // Remove duplicates
        all_cve_ids.sort_unstable();
        all_cve_ids.dedup();

        println!("Total unique CVEs found: {}", all_cve_ids.len());
        Ok(all_cve_ids)
    }

    async fn download_feed(&self, url: &str) -> Result<PathBuf, CollectorError> {
        // Extract filename from URL
        let filename = url.split('/').last().unwrap_or("feed.json.gz");
        let filepath = self.download_dir.join(filename);

        // Check if file already exists to avoid re-downloading
        if filepath.exists() {
            return Ok(filepath);
        }

        // Download the file
        let response = self.client.get(url).send().await?;
        let bytes = response.bytes().await?;

        // Write downloaded bytes to file
        let mut file = File::create(&filepath)?;
        file.write_all(&bytes)?;

        Ok(filepath)
    }

    fn decompress_feed(&self, filepath: &Path) -> Result<Vec<u8>, CollectorError> {
        let file = File::open(filepath)?;
        let mut gz = GzDecoder::new(file);
        let mut buffer = Vec::new();
        gz.read_to_end(&mut buffer)?;

        Ok(buffer)
    }

    fn extract_cve_ids(&self, json_data: &[u8]) -> Result<Vec<String>, CollectorError> {
        let json: Value = serde_json::from_slice(json_data)
            .map_err(|e| CollectorError::NvdFeedError(e.to_string()))?;
        
        // Extract CVE IDs from the JSON feed
        let cve_ids: Vec<String> = json
            .get("CVE_Items")
            .and_then(|items| items.as_array())
            .map(|items_array| {
                items_array
                    .iter()
                    .filter_map(|item| {
                        item.get("cve")
                            .and_then(|cve| cve.get("CVE_data_meta"))
                            .and_then(|meta| meta.get("ID"))
                            .and_then(|id| id.as_str())
                            .map(|s| s.to_string())
                    })
                    .collect()
            })
            .unwrap_or_default();

        Ok(cve_ids)
    }
}

// Modify main.rs to use CLI parsing
#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
    /// Type of NVD feed to download
    #[arg(value_enum, long, default_value_t = FeedType::Recent)]
    feed: FeedType,
}
