use std::collections::BTreeMap;
use std::fs;
use std::path::PathBuf;
use std::time::Duration;

use anyhow::{Context, Result, anyhow};
use chrono::{DateTime, Datelike, SecondsFormat, TimeZone, Utc};
use clap::Parser;
use reqwest::blocking::Client;
use reqwest::header::ACCEPT;
use serde::Serialize;
use serde_json::{Value, json};
use urlencoding::encode;

#[derive(Debug, Parser)]
#[command(about = "Generate DLP compliance reports from AW DLP incidents")]
struct Cli {
    #[arg(long)]
    month: Option<String>,

    #[arg(long)]
    profile: Option<String>,

    #[arg(long, default_value = "152-fz,pci-dss")]
    profiles: String,

    #[arg(long)]
    stdout_json: bool,
}

#[derive(Debug, Clone)]
struct Config {
    aw_api_base: String,
    output_dir: PathBuf,
    base_dir: PathBuf,
    template_152fz: PathBuf,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct ReportStats {
    total_incidents: usize,
    high: usize,
    medium: usize,
    low: usize,
    by_host: BTreeMap<String, usize>,
    channels: BTreeMap<String, usize>,
}

#[derive(Debug, Serialize)]
struct Metadata {
    profile: String,
    period: String,
    generated_at: String,
    aw_api_base: String,
    report_path: String,
    stats: MetadataStats,
}

#[derive(Debug, Serialize)]
struct MetadataStats {
    total_incidents: usize,
    high: usize,
    medium: usize,
    low: usize,
}

fn main() {
    if let Err(err) = run() {
        eprintln!("{err:#}");
        std::process::exit(1);
    }
}

fn run() -> Result<()> {
    let cli = Cli::parse();
    let config = Config::from_env()?;
    let client = Client::builder()
        .timeout(Duration::from_secs(30))
        .no_proxy()
        .build()
        .context("build HTTP client")?;
    let profiles = if let Some(profile) = cli.profile.as_deref() {
        vec![profile.to_string()]
    } else {
        cli.profiles
            .split(',')
            .map(str::trim)
            .filter(|item| !item.is_empty())
            .map(ToOwned::to_owned)
            .collect::<Vec<_>>()
    };
    let mut items = Vec::new();
    for profile in profiles {
        items.push(generate_report(
            &client,
            &config,
            cli.month.as_deref(),
            &profile,
        )?);
    }
    if cli.stdout_json {
        if cli.profile.is_some() && items.len() == 1 {
            println!("{}", serde_json::to_string(&items[0])?);
        } else {
            println!("{}", serde_json::to_string(&json!({"items": items}))?);
        }
    }
    Ok(())
}

impl Config {
    fn from_env() -> Result<Self> {
        let aw_raw = env_string("AW_DLP_AW_API_BASE")
            .or_else(|| env_string("AW_SERVER_URL"))
            .unwrap_or_else(|| "http://127.0.0.1:5600".to_string());
        let aw_api_base = build_aw_api_base(&aw_raw);
        let output_dir = env_path("AW_DLP_COMPLIANCE_REPORT_DIR")
            .unwrap_or_else(|| PathBuf::from("/opt/activitywatch/dlp-compliance/reports"));
        let base_dir = env_path("AW_DLP_COMPLIANCE_BASE_DIR")
            .unwrap_or_else(|| PathBuf::from("/opt/activitywatch/dlp-compliance"));
        let template_152fz = env_path("AW_DLP_COMPLIANCE_TEMPLATE")
            .unwrap_or_else(|| base_dir.join("templates/152-fz-report.html"));
        Ok(Self {
            aw_api_base,
            output_dir,
            base_dir,
            template_152fz,
        })
    }
}

fn generate_report(
    client: &Client,
    config: &Config,
    month: Option<&str>,
    profile: &str,
) -> Result<Metadata> {
    let (start, end, period_label) = period_bounds(month)?;
    let incidents = load_incidents(client, config, start, end)?;
    let stats = build_stats(&incidents);
    fs::create_dir_all(&config.output_dir)
        .with_context(|| format!("create output dir {}", config.output_dir.display()))?;
    let generated_at = Utc::now().to_rfc3339_opts(SecondsFormat::Micros, true);
    let html_path = config
        .output_dir
        .join(format!("{profile}-{period_label}.html"));
    let html = render_html(config, profile, &period_label, &stats, &generated_at)?;
    fs::write(&html_path, html).with_context(|| format!("write {}", html_path.display()))?;
    let metadata = Metadata {
        profile: profile.to_string(),
        period: period_label.clone(),
        generated_at,
        aw_api_base: config.aw_api_base.clone(),
        report_path: html_path.to_string_lossy().to_string(),
        stats: MetadataStats {
            total_incidents: stats.total_incidents,
            high: stats.high,
            medium: stats.medium,
            low: stats.low,
        },
    };
    let json_path = config
        .output_dir
        .join(format!("{profile}-{period_label}.json"));
    fs::write(&json_path, serde_json::to_string_pretty(&metadata)?)
        .with_context(|| format!("write {}", json_path.display()))?;
    Ok(metadata)
}

fn build_aw_api_base(raw_url: &str) -> String {
    let url = raw_url.trim().trim_end_matches('/').to_string();
    if url.ends_with("/api/0") {
        url
    } else {
        format!("{url}/api/0")
    }
}

fn load_incidents(
    client: &Client,
    config: &Config,
    start: DateTime<Utc>,
    end: DateTime<Utc>,
) -> Result<Vec<Value>> {
    let buckets = http_json(client, &format!("{}/buckets", config.aw_api_base))?;
    let Some(map) = buckets.as_object() else {
        return Ok(Vec::new());
    };
    let mut bucket_ids = map
        .keys()
        .filter(|bucket_id| bucket_id.starts_with("aw-dlp-incidents_"))
        .cloned()
        .collect::<Vec<_>>();
    bucket_ids.sort();
    let mut incidents = Vec::new();
    for bucket_id in bucket_ids {
        let encoded = encode(&bucket_id);
        let url = format!("{}/buckets/{encoded}/events?limit=2000", config.aw_api_base);
        let events = http_json(client, &url)?;
        let Some(events) = events.as_array() else {
            continue;
        };
        for event in events {
            let Some(ts) = event
                .get("timestamp")
                .and_then(Value::as_str)
                .and_then(parse_ts)
            else {
                continue;
            };
            if ts < start || ts > end {
                continue;
            }
            incidents.push(event.clone());
        }
    }
    Ok(incidents)
}

fn http_json(client: &Client, url: &str) -> Result<Value> {
    client
        .get(url)
        .header(ACCEPT, "application/json")
        .send()
        .with_context(|| format!("GET {url}"))?
        .error_for_status()
        .with_context(|| format!("GET {url} status"))?
        .json::<Value>()
        .with_context(|| format!("decode JSON from {url}"))
}

fn build_stats(incidents: &[Value]) -> ReportStats {
    let mut by_host = BTreeMap::new();
    let mut channels = BTreeMap::new();
    let mut high = 0;
    let mut medium = 0;
    let mut low = 0;
    for event in incidents {
        let data = event.get("data").and_then(Value::as_object);
        let host = data
            .and_then(|item| item.get("hostname"))
            .and_then(Value::as_str)
            .unwrap_or("unknown")
            .to_string();
        *by_host.entry(host).or_insert(0) += 1;

        match data
            .and_then(|item| item.get("severity"))
            .and_then(Value::as_str)
            .unwrap_or("low")
            .to_lowercase()
            .as_str()
        {
            "high" => high += 1,
            "medium" => medium += 1,
            _ => low += 1,
        }
        let channel = data
            .and_then(|item| item.get("signalType").or_else(|| item.get("source")))
            .and_then(Value::as_str)
            .unwrap_or("unknown")
            .to_string();
        *channels.entry(channel).or_insert(0) += 1;
    }
    ReportStats {
        total_incidents: incidents.len(),
        high,
        medium,
        low,
        by_host: sort_counts_desc(by_host),
        channels: sort_counts_desc(channels),
    }
}

fn sort_counts_desc(values: BTreeMap<String, usize>) -> BTreeMap<String, usize> {
    let mut rows = values.into_iter().collect::<Vec<_>>();
    rows.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
    rows.into_iter().collect()
}

fn render_html(
    config: &Config,
    profile: &str,
    period_label: &str,
    stats: &ReportStats,
    generated_at: &str,
) -> Result<String> {
    let template_path = resolve_template_path(config, profile);
    let template = fs::read_to_string(&template_path)
        .with_context(|| format!("read template {}", template_path.display()))?;
    Ok(template
        .replace("{{PERIOD}}", period_label)
        .replace("{{PROFILE}}", profile)
        .replace("{{GENERATED_AT}}", generated_at)
        .replace("{{TOTAL}}", &stats.total_incidents.to_string())
        .replace("{{HIGH}}", &stats.high.to_string())
        .replace("{{MEDIUM}}", &stats.medium.to_string())
        .replace("{{LOW}}", &stats.low.to_string())
        .replace(
            "{{HOST_TABLE}}",
            &render_table("Инциденты по хостам", &stats.by_host),
        )
        .replace(
            "{{CHANNEL_TABLE}}",
            &render_table("Инциденты по каналам", &stats.channels),
        ))
}

fn resolve_template_path(config: &Config, profile: &str) -> PathBuf {
    if profile == "152-fz" {
        config.template_152fz.clone()
    } else {
        let specific = config
            .base_dir
            .join(format!("templates/{profile}-report.html"));
        if specific.exists() {
            specific
        } else {
            config.template_152fz.clone()
        }
    }
}

fn render_table(title: &str, rows: &BTreeMap<String, usize>) -> String {
    if rows.is_empty() {
        return format!("<h3>{title}</h3><p>Нет данных</p>");
    }
    let body = rows
        .iter()
        .map(|(name, count)| format!("<tr><td>{name}</td><td>{count}</td></tr>"))
        .collect::<String>();
    format!(
        "<h3>{title}</h3><table><thead><tr><th>Параметр</th><th>Значение</th></tr></thead><tbody>{body}</tbody></table>"
    )
}

fn period_bounds(month: Option<&str>) -> Result<(DateTime<Utc>, DateTime<Utc>, String)> {
    let start = if let Some(month) = month {
        let (year, month) = month
            .split_once('-')
            .ok_or_else(|| anyhow!("month must be YYYY-MM"))?;
        Utc.with_ymd_and_hms(year.parse()?, month.parse()?, 1, 0, 0, 0)
            .single()
            .ok_or_else(|| anyhow!("invalid month"))?
    } else {
        let now = Utc::now();
        Utc.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
            .single()
            .ok_or_else(|| anyhow!("invalid current month"))?
    };
    let end = if start.month() == 12 {
        Utc.with_ymd_and_hms(start.year() + 1, 1, 1, 0, 0, 0)
            .single()
            .ok_or_else(|| anyhow!("invalid end month"))?
    } else {
        Utc.with_ymd_and_hms(start.year(), start.month() + 1, 1, 0, 0, 0)
            .single()
            .ok_or_else(|| anyhow!("invalid end month"))?
    };
    Ok((
        start,
        end,
        format!("{:04}-{:02}", start.year(), start.month()),
    ))
}

fn parse_ts(value: &str) -> Option<DateTime<Utc>> {
    DateTime::parse_from_rfc3339(&value.replace('Z', "+00:00"))
        .map(|ts| ts.with_timezone(&Utc))
        .ok()
}

fn env_string(name: &str) -> Option<String> {
    std::env::var(name).ok().filter(|value| !value.is_empty())
}

fn env_path(name: &str) -> Option<PathBuf> {
    env_string(name).map(PathBuf::from)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn aw_api_base_matches_python_contract() {
        assert_eq!(build_aw_api_base("http://x"), "http://x/api/0");
        assert_eq!(build_aw_api_base("http://x/api/0"), "http://x/api/0");
    }

    #[test]
    fn stats_match_expected_counts() {
        let incidents = vec![
            json!({"data":{"hostname":"a","severity":"high","signalType":"usb"}}),
            json!({"data":{"hostname":"a","severity":"medium","source":"clipboard"}}),
            json!({"data":{"hostname":"b","severity":"low"}}),
        ];
        let stats = build_stats(&incidents);
        assert_eq!(stats.total_incidents, 3);
        assert_eq!(stats.high, 1);
        assert_eq!(stats.medium, 1);
        assert_eq!(stats.low, 1);
        assert_eq!(stats.by_host.get("a"), Some(&2));
        assert_eq!(stats.channels.get("unknown"), Some(&1));
    }

    #[test]
    fn month_bounds_are_utc_month() {
        let (start, end, label) = period_bounds(Some("2026-06")).unwrap();
        assert_eq!(label, "2026-06");
        assert_eq!(start.to_rfc3339(), "2026-06-01T00:00:00+00:00");
        assert_eq!(end.to_rfc3339(), "2026-07-01T00:00:00+00:00");
    }
}