use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Write as _;
use std::io::Read as _;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};

use lsp_types as types;
use sha2::{Digest, Sha256};
use shuck_discover::{DiscoveryOptions, FileKind, discover_files_with_status};

use crate::TextDocument;
use crate::lint::generate_diagnostics;
use crate::server::Result;
use crate::session::{
    Client, ClientOptions, ClientSettings, DocumentSnapshot, RequestCancellationToken,
    ShuckSettings, WorkspaceDiagnosticsFeatureOptions, WorkspaceDocumentSnapshotFactory,
    WorkspaceSettingsSnapshot,
};

const PARTIAL_RESULT_BATCH_SIZE: usize = 25;

#[derive(Clone)]
pub(crate) struct WorkspaceDiagnosticContext {
    pub(crate) options: WorkspaceDiagnosticsFeatureOptions,
    pub(crate) global_options: ClientOptions,
    pub(crate) workspace_settings: Vec<WorkspaceSettingsSnapshot>,
    pub(crate) workspace_roots: Vec<PathBuf>,
    pub(crate) settings_workspace_roots: Vec<PathBuf>,
    pub(crate) open_documents: Vec<WorkspaceDiagnosticOpenDocument>,
    pub(crate) snapshot_factory: WorkspaceDocumentSnapshotFactory,
    pub(crate) cache: Arc<WorkspaceDiagnosticCache>,
    pub(crate) cache_generation: u64,
    pub(crate) cancellation: RequestCancellationToken,
}

#[derive(Clone)]
pub(crate) struct WorkspaceDiagnosticOpenDocument {
    pub(crate) uri: types::Url,
    pub(crate) path: PathBuf,
    pub(crate) snapshot: DocumentSnapshot,
}

#[derive(Default)]
pub(crate) struct WorkspaceDiagnosticCache {
    entries: Mutex<BTreeMap<String, CachedWorkspaceDiagnostic>>,
    generation: AtomicU64,
}

#[derive(Clone)]
struct CachedWorkspaceDiagnostic {
    source_hash: [u8; 32],
    result_id: String,
    diagnostics: Vec<types::Diagnostic>,
}

struct WorkDoneGuard {
    client: Client,
    token: Option<types::ProgressToken>,
}

struct ReportSink<'a> {
    client: &'a Client,
    partial_token: Option<&'a types::ProgressToken>,
    buffered: Vec<types::WorkspaceDocumentDiagnosticReport>,
    completed: usize,
}

impl WorkspaceDiagnosticCache {
    pub(crate) fn invalidate_all(&self) {
        self.generation.fetch_add(1, Ordering::AcqRel);
        lock_or_recover(&self.entries).clear();
    }

    pub(crate) fn invalidate_uri(&self, uri: &types::Url) {
        self.generation.fetch_add(1, Ordering::AcqRel);
        lock_or_recover(&self.entries).remove(uri.as_str());
    }

    pub(crate) fn generation(&self) -> u64 {
        self.generation.load(Ordering::Acquire)
    }

    fn get(
        &self,
        uri: &types::Url,
        source_hash: [u8; 32],
        generation: u64,
    ) -> Option<CachedWorkspaceDiagnostic> {
        if self.generation() != generation {
            return None;
        }
        let entries = lock_or_recover(&self.entries);
        if self.generation() != generation {
            return None;
        }
        entries
            .get(uri.as_str())
            .filter(|entry| entry.source_hash == source_hash)
            .cloned()
    }

    fn insert(&self, uri: &types::Url, entry: CachedWorkspaceDiagnostic, generation: u64) {
        if self.generation() != generation {
            return;
        }
        let mut entries = lock_or_recover(&self.entries);
        if self.generation() == generation {
            entries.insert(uri.to_string(), entry);
        }
    }

    fn retain(&self, uris: &BTreeSet<String>, generation: u64) {
        if self.generation() != generation {
            return;
        }
        let mut entries = lock_or_recover(&self.entries);
        if self.generation() == generation {
            entries.retain(|uri, _| uris.contains(uri));
        }
    }
}

impl WorkDoneGuard {
    fn new(client: &Client, token: Option<types::ProgressToken>) -> Self {
        if let Some(token) = token.as_ref() {
            let _ = client.send_notification_value(
                "$/progress",
                serde_json::json!({
                    "token": token,
                    "value": {
                        "kind": "begin",
                        "title": "Shuck workspace diagnostics",
                        "cancellable": true,
                    },
                }),
            );
        }
        Self {
            client: client.clone(),
            token,
        }
    }

    fn report(&self, completed: usize) {
        let Some(token) = self.token.as_ref() else {
            return;
        };
        let _ = self.client.send_notification_value(
            "$/progress",
            serde_json::json!({
                "token": token,
                "value": {
                    "kind": "report",
                    "cancellable": true,
                    "message": format!("Analyzed {completed} shell files"),
                },
            }),
        );
    }
}

impl Drop for WorkDoneGuard {
    fn drop(&mut self) {
        let Some(token) = self.token.as_ref() else {
            return;
        };
        let _ = self.client.send_notification_value(
            "$/progress",
            serde_json::json!({
                "token": token,
                "value": { "kind": "end" },
            }),
        );
    }
}

impl<'a> ReportSink<'a> {
    fn new(client: &'a Client, partial_token: Option<&'a types::ProgressToken>) -> Self {
        Self {
            client,
            partial_token,
            buffered: Vec::new(),
            completed: 0,
        }
    }

    fn completed(&self) -> usize {
        self.completed
    }

    fn push(&mut self, report: types::WorkspaceDocumentDiagnosticReport) -> Result<()> {
        self.buffered.push(report);
        self.completed += 1;
        if self.partial_token.is_some() && self.buffered.len() >= PARTIAL_RESULT_BATCH_SIZE {
            self.flush_partial()?;
        }
        Ok(())
    }

    fn finish(mut self) -> Result<Vec<types::WorkspaceDocumentDiagnosticReport>> {
        if self.partial_token.is_some() {
            self.flush_partial()?;
            return Ok(Vec::new());
        }
        Ok(self.buffered)
    }

    fn flush_partial(&mut self) -> Result<()> {
        let Some(token) = self.partial_token else {
            return Ok(());
        };
        if self.buffered.is_empty() {
            return Ok(());
        }
        self.client.send_notification_value(
            "$/progress",
            serde_json::json!({
                "token": token,
                "value": { "items": self.buffered },
            }),
        )?;
        self.buffered.clear();
        Ok(())
    }
}

pub(crate) fn workspace_diagnostics(
    context: WorkspaceDiagnosticContext,
    client: &Client,
    params: &types::WorkspaceDiagnosticParams,
) -> Result<types::WorkspaceDiagnosticReportResult> {
    let progress = WorkDoneGuard::new(
        client,
        params.work_done_progress_params.work_done_token.clone(),
    );
    if !workspace_diagnostics_have_enabled_scope(&context) || context.cancellation.is_cancelled() {
        return Ok(empty_report());
    }

    let previous = params
        .previous_result_ids
        .iter()
        .map(|previous| (previous.uri.to_string(), previous.value.as_str()))
        .collect::<BTreeMap<_, _>>();
    let mut reports = ReportSink::new(
        client,
        params.partial_result_params.partial_result_token.as_ref(),
    );
    let mut seen = BTreeSet::new();
    let mut open_paths = BTreeSet::new();
    let mut source_bytes = 0usize;
    let mut coverage_complete = true;

    let mut open_documents = context.open_documents.clone();
    open_documents.sort_by(|left, right| left.uri.as_str().cmp(right.uri.as_str()));
    for document in open_documents {
        if reports.completed() >= context.options.max_files {
            coverage_complete = false;
            break;
        }
        if context.cancellation.is_cancelled() {
            coverage_complete = false;
            break;
        }
        let options = workspace_diagnostic_options_for_path(&context, &document.path);
        if !options.enabled || !path_in_workspace(&context, &document.path) {
            continue;
        }
        let document_bytes = document.snapshot.query().document().contents().len();
        if source_bytes.saturating_add(document_bytes) > context.options.max_source_bytes
            || source_bytes.saturating_add(document_bytes) > options.max_source_bytes
        {
            coverage_complete = false;
            continue;
        }
        source_bytes += document_bytes;
        open_paths.insert(canonical_or_original(&document.path));
        let diagnostics = generate_diagnostics(&document.snapshot);
        let result_id = diagnostic_result_id(&diagnostics);
        seen.insert(document.uri.to_string());
        reports.push(document_report(
            document.uri,
            Some(i64::from(document.snapshot.query().document().version())),
            diagnostics,
            result_id,
            &previous,
        ))?;
        progress.report(reports.completed());
    }

    let remaining = context
        .options
        .max_files
        .saturating_sub(reports.completed());
    let mut discovered = BTreeMap::new();
    let mut discovery_entries = 0usize;
    if remaining > 0 && !context.cancellation.is_cancelled() {
        for root in &context.workspace_roots {
            let options = workspace_diagnostic_options_for_path(&context, root);
            if !options.enabled {
                continue;
            }
            if discovered.len() >= remaining {
                coverage_complete = false;
                break;
            }
            let root_remaining = remaining
                .saturating_sub(discovered.len())
                .min(options.max_files);
            let entry_remaining = context
                .options
                .max_entries
                .saturating_sub(discovery_entries)
                .min(options.max_entries);
            if root_remaining == 0 || entry_remaining == 0 {
                coverage_complete = false;
                break;
            }
            let result = match discover_files_with_status(
                std::slice::from_ref(root),
                root,
                &DiscoveryOptions {
                    respect_gitignore: true,
                    parallel: false,
                    use_config_roots: true,
                    max_files: Some(root_remaining),
                    max_entries: Some(entry_remaining),
                    cancellation: Some(context.cancellation.discovery_token()),
                    include_embedded: false,
                    excluded_subtrees: context
                        .workspace_roots
                        .iter()
                        .filter(|other| *other != root && other.starts_with(root))
                        .cloned()
                        .collect(),
                    ..DiscoveryOptions::default()
                },
            ) {
                Ok(files) => files,
                Err(error) => {
                    tracing::warn!(
                        "Failed to discover workspace diagnostic files in {}: {error}",
                        root.display()
                    );
                    coverage_complete = false;
                    continue;
                }
            };
            discovery_entries += result.visited_entries;
            coverage_complete &= result.complete;
            for file in result.files {
                if file.kind != FileKind::Shell
                    || !path_owned_by_root(&context, &file.absolute_path, root)
                    || open_paths.contains(&canonical_or_original(&file.absolute_path))
                {
                    continue;
                }
                discovered.entry(file.absolute_path.clone()).or_insert(file);
            }
        }
    } else if !context.workspace_roots.is_empty() {
        coverage_complete = false;
    }

    for file in discovered.values() {
        if reports.completed() >= context.options.max_files {
            coverage_complete = false;
            break;
        }
        if context.cancellation.is_cancelled() {
            coverage_complete = false;
            break;
        }
        let options = workspace_diagnostic_options_for_path(&context, &file.absolute_path);
        let remaining_source_bytes = context
            .options
            .max_source_bytes
            .saturating_sub(source_bytes)
            .min(options.max_source_bytes.saturating_sub(source_bytes));
        let source = match read_source_bounded(&file.absolute_path, remaining_source_bytes) {
            Ok(Some(source)) => source,
            Ok(None) => {
                coverage_complete = false;
                continue;
            }
            Err(error) => {
                tracing::warn!(
                    "Failed to read workspace diagnostic file {}: {error}",
                    file.absolute_path.display()
                );
                coverage_complete = false;
                continue;
            }
        };
        source_bytes += source.len();
        let Ok(uri) = types::Url::from_file_path(&file.absolute_path) else {
            continue;
        };
        let source_hash: [u8; 32] = Sha256::digest(source.as_bytes()).into();
        let cached = context
            .cache
            .get(&uri, source_hash, context.cache_generation);
        let (diagnostics, result_id) = if let Some(cached) = cached {
            (cached.diagnostics, cached.result_id)
        } else {
            let (settings, client_settings) = settings_for_path(&context, &file.absolute_path);
            let document = Arc::new(TextDocument::new(source, 0));
            let snapshot = context.snapshot_factory.snapshot(
                uri.clone(),
                document,
                Arc::new(settings),
                Arc::new(client_settings),
            );
            let diagnostics = generate_diagnostics(&snapshot);
            let result_id = diagnostic_result_id(&diagnostics);
            context.cache.insert(
                &uri,
                CachedWorkspaceDiagnostic {
                    source_hash,
                    result_id: result_id.clone(),
                    diagnostics: diagnostics.clone(),
                },
                context.cache_generation,
            );
            (diagnostics, result_id)
        };
        seen.insert(uri.to_string());
        reports.push(document_report(
            uri,
            None,
            diagnostics,
            result_id,
            &previous,
        ))?;
        progress.report(reports.completed());
    }

    if coverage_complete && !context.cancellation.is_cancelled() {
        let remaining_reports = context
            .options
            .max_files
            .saturating_sub(reports.completed());
        for previous_uri in previous
            .keys()
            .filter(|uri| !seen.contains(*uri))
            .take(remaining_reports)
        {
            let Ok(uri) = previous_uri.parse::<types::Url>() else {
                continue;
            };
            let Ok(path) = uri.to_file_path() else {
                continue;
            };
            if !path_in_workspace(&context, &path) {
                continue;
            }
            reports.push(document_report(
                uri,
                None,
                Vec::new(),
                diagnostic_result_id(&[]),
                &previous,
            ))?;
        }
    }
    context.cache.retain(&seen, context.cache_generation);

    let reports = reports.finish()?;

    Ok(types::WorkspaceDiagnosticReportResult::Report(
        types::WorkspaceDiagnosticReport { items: reports },
    ))
}

fn read_source_bounded(path: &Path, max_bytes: usize) -> std::io::Result<Option<String>> {
    let metadata = std::fs::metadata(path)?;
    if metadata.len() > max_bytes as u64 {
        return Ok(None);
    }
    let file = std::fs::File::open(path)?;
    let limit = u64::try_from(max_bytes)
        .unwrap_or(u64::MAX)
        .saturating_add(1);
    let mut source = String::new();
    file.take(limit).read_to_string(&mut source)?;
    if source.len() > max_bytes {
        return Ok(None);
    }
    Ok(Some(source))
}

fn empty_report() -> types::WorkspaceDiagnosticReportResult {
    types::WorkspaceDiagnosticReportResult::Report(types::WorkspaceDiagnosticReport {
        items: Vec::new(),
    })
}

fn document_report(
    uri: types::Url,
    version: Option<i64>,
    diagnostics: Vec<types::Diagnostic>,
    result_id: String,
    previous: &BTreeMap<String, &str>,
) -> types::WorkspaceDocumentDiagnosticReport {
    if previous.get(uri.as_str()).copied() == Some(result_id.as_str()) {
        return types::WorkspaceDocumentDiagnosticReport::Unchanged(
            types::WorkspaceUnchangedDocumentDiagnosticReport {
                uri,
                version,
                unchanged_document_diagnostic_report: types::UnchangedDocumentDiagnosticReport {
                    result_id,
                },
            },
        );
    }
    types::WorkspaceDocumentDiagnosticReport::Full(types::WorkspaceFullDocumentDiagnosticReport {
        uri,
        version,
        full_document_diagnostic_report: types::FullDocumentDiagnosticReport {
            result_id: Some(result_id),
            items: diagnostics,
        },
    })
}

fn diagnostic_result_id(diagnostics: &[types::Diagnostic]) -> String {
    let serialized = serde_json::to_vec(diagnostics).unwrap_or_default();
    let digest = Sha256::digest(&serialized);
    let mut result = String::with_capacity(32);
    for byte in &digest[..16] {
        let _ = write!(result, "{byte:02x}");
    }
    result
}

fn settings_for_path(
    context: &WorkspaceDiagnosticContext,
    path: &Path,
) -> (ShuckSettings, ClientSettings) {
    let workspace_options =
        workspace_settings_for_path(context, path).and_then(|workspace| workspace.options.as_ref());
    if let Some(workspace_options) = workspace_options {
        let layers = [&context.global_options, workspace_options];
        return (
            ShuckSettings::resolve(path.into(), &context.settings_workspace_roots, &layers),
            ClientSettings::from_layered_options(&layers),
        );
    }
    let layers = [&context.global_options];
    (
        ShuckSettings::resolve(path.into(), &context.settings_workspace_roots, &layers),
        ClientSettings::from_layered_options(&layers),
    )
}

fn workspace_diagnostic_options_for_path(
    context: &WorkspaceDiagnosticContext,
    path: &Path,
) -> WorkspaceDiagnosticsFeatureOptions {
    workspace_settings_for_path(context, path)
        .and_then(|workspace| workspace.options.as_ref())
        .map(|options| {
            options
                .server
                .workspace_diagnostics_layered_over(context.options)
        })
        .unwrap_or(context.options)
}

fn workspace_settings_for_path<'a>(
    context: &'a WorkspaceDiagnosticContext,
    path: &Path,
) -> Option<&'a WorkspaceSettingsSnapshot> {
    context
        .workspace_settings
        .iter()
        .filter_map(|workspace| {
            workspace_root_match_len(path, workspace).map(|len| (workspace, len))
        })
        .max_by_key(|(_, len)| *len)
        .map(|(workspace, _)| workspace)
}

fn workspace_root_match_len(path: &Path, workspace: &WorkspaceSettingsSnapshot) -> Option<usize> {
    [Some(&workspace.root), workspace.canonical_root.as_ref()]
        .into_iter()
        .flatten()
        .filter(|root| path.starts_with(root))
        .map(|root| root.components().count())
        .max()
}

fn workspace_diagnostics_have_enabled_scope(context: &WorkspaceDiagnosticContext) -> bool {
    context
        .workspace_roots
        .iter()
        .any(|root| workspace_diagnostic_options_for_path(context, root).enabled)
}

fn path_in_workspace(context: &WorkspaceDiagnosticContext, path: &Path) -> bool {
    if workspace_settings_for_path(context, path).is_some() {
        return true;
    }

    let canonical_path = crate::workspace_functions::canonical_path(path);
    workspace_settings_for_path(context, &canonical_path).is_some()
        || context
            .workspace_roots
            .iter()
            .any(|root| path_starts_with_root(path, root))
}

fn path_starts_with_root(path: &Path, root: &Path) -> bool {
    path.starts_with(root)
        || crate::workspace_functions::canonical_path(path)
            .starts_with(crate::workspace_functions::canonical_path(root))
}

fn path_owned_by_root(context: &WorkspaceDiagnosticContext, path: &Path, root: &Path) -> bool {
    workspace_settings_for_path(context, path)
        .map(|workspace| workspace.root == root)
        .unwrap_or(true)
}

fn canonical_or_original(path: &Path) -> PathBuf {
    path.canonicalize().unwrap_or_else(|_| path.to_path_buf())
}

fn lock_or_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    match mutex.lock() {
        Ok(guard) => guard,
        Err(poisoned) => poisoned.into_inner(),
    }
}

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

    #[test]
    fn stale_request_cannot_repopulate_an_invalidated_cache() {
        let cache = WorkspaceDiagnosticCache::default();
        let uri = types::Url::parse("file:///workspace/script.sh").unwrap();
        let stale_generation = cache.generation();
        cache.invalidate_all();
        cache.insert(
            &uri,
            CachedWorkspaceDiagnostic {
                source_hash: [7; 32],
                result_id: "stale".to_owned(),
                diagnostics: Vec::new(),
            },
            stale_generation,
        );

        assert!(cache.get(&uri, [7; 32], cache.generation()).is_none());
    }

    #[test]
    fn bounded_reader_rejects_a_file_before_loading_past_the_budget() {
        let directory = tempfile::tempdir().unwrap();
        let path = directory.path().join("large.sh");
        std::fs::write(&path, vec![b'x'; 4096]).unwrap();

        assert!(read_source_bounded(&path, 16).unwrap().is_none());
    }

    #[test]
    fn canonical_deleted_path_remains_under_workspace_root() {
        let directory = tempfile::tempdir().unwrap();
        let path = directory.path().join("deleted.sh");
        std::fs::write(&path, "unused=1\n").unwrap();
        let canonical_path = crate::workspace_functions::canonical_path(&path);
        std::fs::remove_file(&path).unwrap();

        assert!(path_starts_with_root(&canonical_path, directory.path()));
    }
}
