//! Diagnostics jobs (port of `diagnosticsService.js`): one running job, two tools at a time,
//! at most two IP families, five jobs kept, 500 progress events, time limit 5 min (or the
//! ping budget, max 30 min), cancellation of the whole job or of a single tool.

use crate::core::diagnostics::common::{error_result, normalize_options, now_ms, parse_target, DiagError, Options, PING_DURATION_LIMIT_MS};
use crate::core::diagnostics::history::{export_report, scrubbed, DiagnosticsHistory};
use crate::core::diagnostics::http::HttpDiagnostics;
use crate::core::diagnostics::{dns, process, socket};
use parking_lot::Mutex;
use serde_json::{json, Map, Value};
use std::collections::HashMap;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tokio_util::sync::CancellationToken;

pub type ProgressSink = Arc<dyn Fn(Value) + Send + Sync>;

struct Job {
    report: Value,
    cancel: CancellationToken,
    tools: HashMap<String, CancellationToken>,
    cancelled_tools: Vec<String>,
}

pub struct DiagnosticsService {
    pub history: DiagnosticsHistory,
    http: HttpDiagnostics,
    jobs: Mutex<Vec<(String, Arc<Mutex<Job>>)>>,
    resolvers: Mutex<Vec<(u64, CancellationToken)>>,
    resolver_seq: std::sync::atomic::AtomicU64,
    sink: ProgressSink,
}

struct ResolverGuard<'a> {
    service: &'a DiagnosticsService,
    ticket: u64,
    cancel: CancellationToken,
}

impl Drop for ResolverGuard<'_> {
    fn drop(&mut self) {
        self.cancel.cancel();
        self.service.resolvers.lock().retain(|(ticket, _)| *ticket != self.ticket);
    }
}

pub fn job_time_limit(options: &Options) -> u64 {
    if options.continuous {
        return PING_DURATION_LIMIT_MS;
    }
    let ping = if options.tools.iter().any(|t| t == "ping") { options.count * options.repetitions * options.interval_ms + options.ping_timeout_ms.unwrap_or(options.timeout_ms) + 60_000 } else { 0 };
    PING_DURATION_LIMIT_MS.min((5 * 60 * 1000u64).max(ping))
}

fn target_for_address(options: &Options, address: &str) -> Result<Options, DiagError> {
    let mut url = url::Url::parse(&options.target.url).map_err(|_| DiagError::new("INVALID_URL"))?;
    let host = if address.parse::<IpAddr>().map(|ip| ip.is_ipv6()).unwrap_or(false) { format!("[{address}]") } else { address.to_string() };
    url.set_host(Some(&host)).map_err(|_| DiagError::new("INVALID_URL"))?;
    Ok(Options { target: parse_target(url.as_str())?, ..options.clone() })
}

fn same_ip(a: &str, b: &str) -> bool {
    match (a.parse::<IpAddr>(), b.parse::<IpAddr>()) {
        (Ok(x), Ok(y)) => x == y,
        _ => false,
    }
}

fn merge(mut base: Value, extra: &Value) -> Value {
    if let (Some(target), Some(fields)) = (base.as_object_mut(), extra.as_object()) {
        for (k, v) in fields {
            target.insert(k.clone(), v.clone());
        }
    }
    base
}

impl DiagnosticsService {
    pub fn new(storage_root: &Path, resources: PathBuf, sink: ProgressSink) -> Self {
        Self { history: DiagnosticsHistory::new(Some(storage_root)), http: HttpDiagnostics { resources, ca_path: None }, jobs: Mutex::new(Vec::new()), resolvers: Mutex::new(Vec::new()), resolver_seq: Default::default(), sink }
    }

    fn reserve_resolver(&self) -> Result<ResolverGuard<'_>, DiagError> {
        let mut resolvers = self.resolvers.lock();
        if resolvers.len() >= 2 {
            return Err(DiagError::new("DIAGNOSTICS_BUSY"));
        }
        let ticket = self.resolver_seq.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let cancel = CancellationToken::new();
        resolvers.push((ticket, cancel.clone()));
        Ok(ResolverGuard { service: self, ticket, cancel })
    }

    pub async fn resolve(&self, source: &Value) -> Result<Value, DiagError> {
        let mut input = source.clone();
        input["tools"] = json!(["dns"]);
        let options = normalize_options(&input)?;
        let reservation = self.reserve_resolver()?;
        let cancel = reservation.cancel.clone();
        let result = async {
            if let Some(manual) = &options.manual_addresses {
                let mut target = options.target.to_json();
                target["type"] = json!("ipv4+ipv6");
                return Ok(json!({ "target": target, "addresses": manual.iter().map(|(a, f)| json!({ "address": a, "family": f })).collect::<Vec<_>>() }));
            }
            let addresses = tokio::time::timeout(Duration::from_secs(30), dns::resolve(&options.target.host, &options, &cancel)).await.map_err(|_| DiagError::new("TIMEOUT"))??;
            Ok(json!({ "target": options.target.to_json(), "addresses": addresses }))
        }
        .await;
        drop(reservation);
        result
    }

    pub fn start(self: &Arc<Self>, source: &Value) -> Result<Value, DiagError> {
        let options = normalize_options(source)?;
        let mut id_bytes = [0u8; 16];
        rand::RngCore::fill_bytes(&mut rand::rngs::OsRng, &mut id_bytes);
        let id = hex::encode(id_bytes);
        let report = json!({ "id": id, "createdAt": chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true), "status": "running", "options": options.json, "results": {}, "progress": [], "addresses": [] });
        let job = Arc::new(Mutex::new(Job { report: report.clone(), cancel: CancellationToken::new(), tools: HashMap::new(), cancelled_tools: Vec::new() }));
        {
            let mut jobs = self.jobs.lock();
            if jobs.iter().any(|(_, job)| job.lock().report["status"] == "running") {
                return Err(DiagError::new("DIAGNOSTICS_BUSY"));
            }
            jobs.push((id.clone(), job.clone()));
            while jobs.len() > 5 {
                jobs.remove(0);
            }
        }
        let service = self.clone();
        tokio::spawn(async move {
            let outcome = service.run(&job, options).await;
            if let Err(error) = outcome {
                service.finish(&job, Some(error_result(&error)));
            }
        });
        Ok(json!({ "id": id, "report": scrubbed(&report) }))
    }

    pub fn cancel(&self, id: &str, tool: Option<&str>) -> Result<Value, DiagError> {
        let job = self.jobs.lock().iter().find(|(jid, _)| jid == id).map(|(_, j)| j.clone()).ok_or_else(|| DiagError::new("JOB_NOT_FOUND"))?;
        let mut guard = job.lock();
        match tool.filter(|t| !t.is_empty()) {
            Some(tool) => {
                if !guard.report["options"]["tools"].as_array().map(|t| t.iter().any(|x| x == tool)).unwrap_or(false) {
                    return Err(DiagError::new("INVALID_TOOLS"));
                }
                guard.cancelled_tools.push(tool.to_string());
                if let Some(token) = guard.tools.get(tool) {
                    token.cancel();
                }
            }
            None => guard.cancel.cancel(),
        }
        Ok(json!({ "cancelling": true }))
    }

    pub fn stop_all(&self) {
        for (_, job) in self.jobs.lock().iter() {
            let guard = job.lock();
            if guard.report["status"] == "running" {
                guard.cancel.cancel();
            }
        }
        for (_, token) in self.resolvers.lock().iter() {
            token.cancel();
        }
    }

    pub fn get(&self, id: &str) -> Result<Value, DiagError> {
        if let Some(job) = self.jobs.lock().iter().find(|(jid, _)| jid == id).map(|(_, j)| j.clone()) {
            return Ok(scrubbed(&job.lock().report));
        }
        self.history.get(id)
    }

    pub fn export(&self, id: &str, format: &str) -> Result<String, DiagError> {
        export_report(&self.get(id)?, format)
    }

    fn emit(&self, job: &Arc<Mutex<Job>>, event: Value) {
        let mut guard = job.lock();
        let id = guard.report["id"].clone();
        let kind = event["kind"].as_str().unwrap_or("").to_string();
        if kind != "snapshot" && kind != "complete" {
            if let Some(progress) = guard.report["progress"].as_array_mut() {
                progress.push(scrubbed(&event));
                if progress.len() > 500 {
                    progress.remove(0);
                }
            }
        }
        drop(guard);
        (self.sink)(merge(json!({ "jobId": id }), &event));
    }

    fn finish(&self, job: &Arc<Mutex<Job>>, failure: Option<Value>) {
        let report = {
            let mut guard = job.lock();
            if guard.report["status"] != "running" {
                return;
            }
            let aborted = guard.cancel.is_cancelled();
            guard.report["status"] = json!(if aborted { "cancelled" } else if failure.is_some() { "failed" } else { "completed" });
            guard.report["finishedAt"] = json!(chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true));
            if let Some(failure) = failure {
                guard.report["error"] = failure;
            }
            if let Some(reason) = guard.report.get("_abortReason").cloned() {
                guard.report["error"] = reason;
                if let Some(o) = guard.report.as_object_mut() {
                    o.remove("_abortReason");
                }
            }
            guard.report.clone()
        };
        if let Err(error) = self.history.add(&report) {
            job.lock().report["historyError"] = json!(error.code);
        }
        let snapshot = scrubbed(&job.lock().report);
        self.emit(job, json!({ "kind": "complete", "report": snapshot }));
    }

    async fn address_run(&self, job: &Arc<Mutex<Job>>, tool: &str, item: &Value, options: &Options, cancel: &CancellationToken, progress: ProgressSink) -> Value {
        let address = item["address"].as_str().unwrap_or("").to_string();
        let options = if options.manual_addresses.is_some() { target_for_address(options, &address).unwrap_or_else(|_| options.clone()) } else { options.clone() };
        let outcome: Result<Value, DiagError> = async {
            if cancel.is_cancelled() {
                return Err(DiagError::new("CANCELLED"));
            }
            let mut result = match tool {
                "http" => self.http.run(&options.target, &address, &options, cancel, progress.clone()).await?,
                "compression" => self.http.compression(&options.target, &address, &options, cancel, progress.clone()).await?,
                "tls" => socket::run_tls(&options.target, &address, &options, cancel, progress.clone(), None).await?,
                "websocket" => socket::websocket_probe(&options.target, &address, &options, cancel, None).await?,
                "ping" => process::ping(&address, &options, cancel, progress.clone()).await?,
                "traceroute" => process::traceroute(&address, &options, cancel, progress.clone()).await?,
                "tcp" => socket::run_tcp(&address, &options, cancel, progress.clone()).await?,
                _ => return Err(DiagError::new("INVALID_TOOLS")),
            };
            if tool == "traceroute" {
                let mut hops: Vec<String> = Vec::new();
                for route in result["routes"].as_array().cloned().unwrap_or_default() {
                    for hop in route["hops"].as_array().cloned().unwrap_or_default() {
                        for a in hop["addresses"].as_array().cloned().unwrap_or_default() {
                            if let Some(a) = a.as_str() {
                                if !hops.contains(&a.to_string()) {
                                    hops.push(a.to_string());
                                }
                            }
                        }
                    }
                }
                hops.truncate(options.max_hops as usize);
                let mut reverse = Map::new();
                for hop in hops {
                    if cancel.is_cancelled() {
                        return Err(DiagError::new("CANCELLED"));
                    }
                    let names = match dns::reverse_name(&hop) {
                        Ok(name) => match dns::query(&name, "PTR", &options, cancel).await {
                            Ok(answer) => answer["records"].as_array().cloned().unwrap_or_default().into_iter().filter(|r| r["type"] == "PTR").map(|r| r["value"].clone()).collect(),
                            Err(error) => {
                                if cancel.is_cancelled() {
                                    return Err(error);
                                }
                                Vec::new()
                            }
                        },
                        Err(_) => Vec::new(),
                    };
                    reverse.insert(hop, Value::Array(names));
                }
                result["reverseDns"] = Value::Object(reverse);
            }
            Ok(result)
        }
        .await;
        let base = json!({ "address": address, "family": item["family"] });
        match outcome {
            Ok(result) => merge(base, &result),
            Err(error) => {
                let mut partial = merge(base, &error_result(&error));
                let report = job.lock().report.clone();
                let progress_list = report["progress"].as_array().cloned().unwrap_or_default();
                match tool {
                    "ping" => {
                        if let Some(latest) = report["livePing"].get(&address) {
                            for key in ["sent", "received", "lossPercent", "statistics"] {
                                partial[key] = latest[key].clone();
                            }
                        }
                        partial["packets"] = Value::Array(progress_list.into_iter().filter(|e| e["kind"] == "ping-packet" && e["address"] == address.as_str()).collect());
                    }
                    "traceroute" => {
                        partial["routes"] = json!([{ "hops": progress_list.into_iter().filter(|e| e["kind"] == "trace-hop" && e["address"] == address.as_str()).collect::<Vec<_>>() }]);
                    }
                    "http" => {
                        partial["comparisons"] = Value::Array(
                            options
                                .protocols
                                .iter()
                                .map(|p| json!({ "protocol": p, "address": address, "samples": progress_list.iter().filter(|e| e["kind"] == "http-sample" && e["address"] == address.as_str() && e["protocol"] == p.as_str()).cloned().collect::<Vec<_>>() }))
                                .collect(),
                        );
                    }
                    _ => {}
                }
                partial
            }
        }
    }

    async fn run(self: &Arc<Self>, job: &Arc<Mutex<Job>>, options: Options) -> Result<(), DiagError> {
        let cancel = job.lock().cancel.clone();
        let deadline_cancel = cancel.clone();
        let deadline_job = job.clone();
        let limit = job_time_limit(&options);
        let deadline = tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(limit)).await;
            deadline_job.lock().report["_abortReason"] = error_result(&DiagError::new("JOB_TIME_LIMIT"));
            deadline_cancel.cancel();
        });
        let outcome = async {
            if cancel.is_cancelled() {
                return Err(DiagError::new("CANCELLED"));
            }
            let started = now_ms();
            let mut addresses: Vec<Value> = match &options.manual_addresses {
                Some(manual) => manual.iter().map(|(a, f)| json!({ "address": a, "family": f, "source": "manual" })).collect(),
                None => match dns::resolve(&options.target.host, &options, &cancel).await {
                    Ok(list) => list,
                    Err(error) => {
                        if cancel.is_cancelled() {
                            return Err(error);
                        }
                        job.lock().report["resolutionError"] = error_result(&error);
                        Vec::new()
                    }
                },
            };
            let resolved = addresses.clone();
            let resolution_ms = now_ms() - started;
            if !options.selected_ip.is_empty() {
                addresses.retain(|a| same_ip(a["address"].as_str().unwrap_or(""), &options.selected_ip));
                if addresses.is_empty() {
                    return Err(DiagError::new("SELECTED_IP_NOT_RESOLVED"));
                }
            }
            let family = |a: &Value| a["family"].as_u64().unwrap_or(0);
            match options.mode.as_str() {
                "ipv4" => addresses.retain(|a| family(a) == 4),
                "ipv6" => addresses.retain(|a| family(a) == 6),
                "auto" if options.manual_addresses.is_none() => {
                    addresses = addresses.iter().find(|a| family(a) == 4).or(addresses.first()).cloned().into_iter().collect();
                }
                "both" => {
                    let v4 = addresses.iter().find(|a| family(a) == 4).cloned();
                    let v6 = addresses.iter().find(|a| family(a) == 6).cloned();
                    addresses = v4.into_iter().chain(v6).collect();
                }
                _ => {}
            }
            let resolution_error = {
                let mut guard = job.lock();
                guard.report["addresses"] = json!(addresses);
                guard.report["resolvedAddresses"] = json!(resolved);
                guard.report["resolutionMs"] = json!(resolution_ms);
                guard.report["resolutionError"].clone()
            };
            self.emit(job, json!({ "kind": "resolved", "addresses": addresses, "resolvedAddresses": resolved, "resolutionMs": resolution_ms, "error": resolution_error }));
            let queue = Arc::new(Mutex::new(std::collections::VecDeque::from(options.tools.clone())));
            let workers = (0..options.tools.len().min(2)).map(|_| {
                let queue = queue.clone();
                let service = self.clone();
                let job = job.clone();
                let options = options.clone();
                let cancel = cancel.clone();
                let addresses = addresses.clone();
                async move {
                    loop {
                        if cancel.is_cancelled() {
                            break;
                        }
                        let Some(tool) = queue.lock().pop_front() else { break };
                        let sub = cancel.child_token();
                        {
                            let mut guard = job.lock();
                            if guard.cancelled_tools.contains(&tool) {
                                sub.cancel();
                            }
                            guard.tools.insert(tool.clone(), sub.clone());
                        }
                        service.emit(&job, json!({ "kind": "tool-start", "tool": tool }));
                        let progress_job = job.clone();
                        let progress_service = service.clone();
                        let progress_tool = tool.clone();
                        let progress: ProgressSink = Arc::new(move |event: Value| {
                            if event["kind"] == "ping-packet" {
                                let address = event["address"].as_str().unwrap_or("").to_string();
                                let mut guard = progress_job.lock();
                                if !guard.report["livePing"].is_object() {
                                    guard.report["livePing"] = json!({});
                                }
                                guard.report["livePing"][address] = event.clone();
                            }
                            progress_service.emit(&progress_job, merge(json!({ "tool": progress_tool }), &event));
                        });
                        let result = if sub.is_cancelled() {
                            json!([error_result(&DiagError::new("CANCELLED"))])
                        } else if tool == "dns" {
                            match dns::run(&options.target.host, &options, &sub, &*progress).await {
                                Ok(result) => result,
                                Err(error) => json!([error_result(&error)]),
                            }
                        } else if addresses.is_empty() {
                            json!([{ "status": "not-tested", "code": "NO_ADDRESS" }])
                        } else {
                            let futures = addresses.iter().map(|item| service.address_run(&job, &tool, item, &options, &sub, progress.clone()));
                            Value::Array(futures::future::join_all(futures).await)
                        };
                        job.lock().report["results"][&tool] = result.clone();
                        service.emit(&job, json!({ "kind": "tool-complete", "tool": tool, "result": scrubbed(&result) }));
                        job.lock().tools.remove(&tool);
                    }
                }
            });
            futures::future::join_all(workers).await;
            Ok(())
        }
        .await;
        deadline.abort();
        match outcome {
            Ok(()) => self.finish(job, None),
            Err(error) => self.finish(job, Some(error_result(&error))),
        }
        Ok(())
    }
}

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

    #[test]
    fn concurrent_resolvers_reserve_two_slots_and_release_on_drop() {
        let root = tempfile::tempdir().unwrap();
        let service = DiagnosticsService::new(root.path(), PathBuf::new(), Arc::new(|_| {}));
        let barrier = Barrier::new(16);
        std::thread::scope(|scope| {
            let handles: Vec<_> = (0..16).map(|_| scope.spawn(|| {
                let reservation = service.reserve_resolver();
                barrier.wait();
                let accepted = reservation.is_ok();
                barrier.wait();
                accepted
            })).collect();
            assert_eq!(handles.into_iter().map(|h| h.join().unwrap()).filter(|accepted| *accepted).count(), 2);
        });
        assert!(service.resolvers.lock().is_empty());
        assert!(service.reserve_resolver().is_ok());
    }

    #[tokio::test]
    async fn concurrent_starts_reserve_one_job() {
        // The current-thread runtime does not run spawned jobs until all starters return.
        let root = tempfile::tempdir().unwrap();
        let service = Arc::new(DiagnosticsService::new(root.path(), PathBuf::new(), Arc::new(|_| {})));
        let runtime = tokio::runtime::Handle::current();
        let barrier = Barrier::new(16);
        std::thread::scope(|scope| {
            let handles: Vec<_> = (0..16).map(|_| scope.spawn(|| {
                let _entered = runtime.enter();
                barrier.wait();
                service.start(&json!({ "target": "example.com", "tools": ["dns"] })).is_ok()
            })).collect();
            assert_eq!(handles.into_iter().map(|h| h.join().unwrap()).filter(|accepted| *accepted).count(), 1);
        });
        service.stop_all();
    }
}

#[cfg(test)]
#[path = "../../../tests/support/diagnostics_safety.rs"]
mod safety_tests;