#![cfg_attr(target_family = "wasm", allow(dead_code))]

use std::borrow::Cow;
use std::fs::File;
use std::io;
use std::io::Write;
use std::ops::{Add, AddAssign};
use std::path::Path;

use anyhow::{Context, Result};
use colored::Colorize;
use log::{debug, warn};
use ruff_db::diagnostic::Diagnostic;
use ruff_linter::codes::Rule;
use ruff_linter::linter::{FixTable, FixerResult, LinterResult, ParseSource, lint_fix, lint_only};
use ruff_linter::package::PackageRoot;
use ruff_linter::settings::types::UnsafeFixes;
use ruff_linter::settings::{LinterSettings, flags};
use ruff_linter::source_kind::{SourceError, SourceKind, SourceKindDiff};
use ruff_linter::toml::{TomlFixerResult, lint_fix_toml, lint_toml};
use ruff_linter::{IOError, Violation, fs};
use ruff_notebook::{NotebookError, NotebookIndex};
use ruff_python_ast::{SourceType, TomlSourceType};
use ruff_source_file::SourceFileBuilder;
use ruff_text_size::TextRange;
use ruff_workspace::Settings;
use rustc_hash::FxHashMap;

use crate::cache::{Cache, FileCache, FileCacheKey};

/// A collection of [`Diagnostic`]s and additional information needed to render them.
///
/// Note that `notebook_indexes` may be empty if there are no diagnostics because the
/// `NotebookIndex` isn't cached in this case. This isn't a problem for any current uses as of
/// 2025-08-12, which are all related to diagnostic rendering, but could be surprising if used
/// differently in the future.
#[derive(Debug, Default, PartialEq)]
pub(crate) struct Diagnostics {
    pub(crate) inner: Vec<Diagnostic>,
    pub(crate) fixed: FixMap,
    pub(crate) notebook_indexes: FxHashMap<String, NotebookIndex>,
}

impl Diagnostics {
    pub(crate) fn new(
        diagnostics: Vec<Diagnostic>,
        notebook_indexes: FxHashMap<String, NotebookIndex>,
    ) -> Self {
        Self {
            inner: diagnostics,
            fixed: FixMap::default(),
            notebook_indexes,
        }
    }

    /// Generate [`Diagnostics`] based on a [`SourceError`].
    fn from_source_error(
        err: &SourceError,
        path: Option<&Path>,
        settings: &LinterSettings,
    ) -> Self {
        match err {
            // IO errors.
            SourceError::Io(_)
            | SourceError::Notebook(NotebookError::Io(_) | NotebookError::Json(_)) => {
                if settings.rules.enabled(Rule::IOError) {
                    let name = path.map_or_else(|| "-".into(), Path::to_string_lossy);
                    let source_file = SourceFileBuilder::new(name, "").finish();
                    Self::new(
                        vec![
                            IOError {
                                message: err.to_string(),
                            }
                            .into_diagnostic(TextRange::default(), &source_file),
                        ],
                        FxHashMap::default(),
                    )
                } else {
                    match path {
                        Some(path) => {
                            warn!(
                                "{}{}{} {err}",
                                "Failed to lint ".bold(),
                                fs::relativize_path(path).bold(),
                                ":".bold()
                            );
                        }
                        None => {
                            warn!("{}{} {err}", "Failed to lint".bold(), ":".bold());
                        }
                    }

                    Self::default()
                }
            }
            // Syntax errors.
            SourceError::Notebook(
                NotebookError::InvalidJson(_)
                | NotebookError::InvalidSchema(_)
                | NotebookError::InvalidFormat(_),
            ) => {
                let name = path.map_or_else(|| "-".into(), Path::to_string_lossy);
                let dummy = SourceFileBuilder::new(name, "").finish();
                Self::new(
                    vec![Diagnostic::invalid_syntax(dummy, err, TextRange::default())],
                    FxHashMap::default(),
                )
            }
        }
    }
}

impl Add for Diagnostics {
    type Output = Diagnostics;

    fn add(mut self, other: Self) -> Self::Output {
        self += other;
        self
    }
}

impl AddAssign for Diagnostics {
    fn add_assign(&mut self, other: Self) {
        self.inner.extend(other.inner);
        self.fixed += other.fixed;
        self.notebook_indexes.extend(other.notebook_indexes);
    }
}

/// A collection of fixes indexed by file path.
#[derive(Debug, Default, PartialEq)]
pub(crate) struct FixMap(FxHashMap<String, FixTable>);

impl FixMap {
    /// Returns `true` if there are no fixes in the map.
    pub(crate) fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Returns an iterator over the fixes in the map, along with the file path.
    pub(crate) fn iter(&self) -> impl Iterator<Item = (&String, &FixTable)> {
        self.0.iter()
    }

    /// Returns an iterator over the fixes in the map.
    pub(crate) fn values(&self) -> impl Iterator<Item = &FixTable> {
        self.0.values()
    }
}

impl FromIterator<(String, FixTable)> for FixMap {
    fn from_iter<T: IntoIterator<Item = (String, FixTable)>>(iter: T) -> Self {
        Self(
            iter.into_iter()
                .filter(|(_, fixes)| !fixes.is_empty())
                .collect(),
        )
    }
}

impl AddAssign for FixMap {
    fn add_assign(&mut self, rhs: Self) {
        #[expect(
            clippy::iter_over_hash_type,
            reason = "fix counts are merged independently for each filename"
        )]
        for (filename, fixed) in rhs.0 {
            if fixed.is_empty() {
                continue;
            }
            let fixed_in_file = self.0.entry(filename).or_default();
            for (id, code, count) in fixed.iter() {
                if count > 0 {
                    *fixed_in_file.entry(id).or_default(code) += count;
                }
            }
        }
    }
}

/// Lint the source code at the given `Path`.
pub(crate) fn lint_path(
    path: &Path,
    package: Option<PackageRoot<'_>>,
    settings: &LinterSettings,
    cache: Option<&Cache>,
    noqa: flags::Noqa,
    fix_mode: flags::FixMode,
    unsafe_fixes: UnsafeFixes,
) -> Result<Diagnostics> {
    // Check the cache.
    let caching = match cache {
        Some(cache) if noqa.is_enabled() => {
            let relative_path = cache
                .relative_path(path)
                .expect("wrong package cache for file");

            let cache_key = FileCacheKey::from_path(path).context("Failed to create cache key")?;
            let cached_diagnostics = cache
                .get(relative_path, &cache_key)
                .is_some_and(FileCache::linted);
            if cached_diagnostics {
                return Ok(Diagnostics::default());
            }

            // Stash the file metadata for later so when we update the cache it reflects the prerun
            // information
            Some((cache, relative_path, cache_key))
        }
        _ => None,
    };

    debug!("Checking: {}", path.display());

    let source_type = match settings.extension.get_source_type(path) {
        SourceType::Toml(source_type @ (TomlSourceType::Pyproject | TomlSourceType::Ruff)) => {
            let (diagnostics, fixed) = if settings
                .rules
                .iter_enabled()
                .any(|rule_code| rule_code.lint_source().is_toml())
            {
                let contents = match std::fs::read_to_string(path).map_err(SourceError::from) {
                    Ok(contents) => contents,
                    Err(err) => {
                        return Ok(Diagnostics::from_source_error(&err, Some(path), settings));
                    }
                };
                if matches!(fix_mode, flags::FixMode::Apply | flags::FixMode::Diff) {
                    let TomlFixerResult {
                        diagnostics,
                        transformed,
                        fixed,
                    } = lint_fix_toml(path, &contents, settings, source_type, unsafe_fixes);

                    if !fixed.is_empty() {
                        match fix_mode {
                            flags::FixMode::Apply => {
                                File::create(path)?.write_all(transformed.as_bytes())?;
                            }
                            flags::FixMode::Diff => {
                                write!(
                                    &mut io::stdout().lock(),
                                    "{}",
                                    SourceKindDiff::from_text(
                                        &contents,
                                        transformed.as_ref(),
                                        Some(path),
                                    )
                                )?;
                            }
                            flags::FixMode::Generate => {}
                        }
                    }

                    (diagnostics, fixed)
                } else {
                    (
                        lint_toml(path, &contents, settings, source_type),
                        FixTable::default(),
                    )
                }
            } else {
                (vec![], FixTable::default())
            };
            return Ok(Diagnostics {
                inner: diagnostics,
                fixed: FixMap::from_iter([(fs::relativize_path(path), fixed)]),
                notebook_indexes: FxHashMap::default(),
            });
        }
        SourceType::Toml(_) | SourceType::Markdown => return Ok(Diagnostics::default()),
        SourceType::Python(source_type) => source_type,
    };

    // Extract the sources from the file.
    let source_kind = match SourceKind::from_path(path, SourceType::Python(source_type)) {
        Ok(Some(source_kind)) => match source_kind {
            SourceKind::Markdown(_) => return Ok(Diagnostics::default()), // skip linting markdown
            _ => source_kind,
        },
        Ok(None) => return Ok(Diagnostics::default()),
        Err(err) => {
            return Ok(Diagnostics::from_source_error(&err, Some(path), settings));
        }
    };

    // Lint the file.
    let (result, transformed, fixed) =
        if matches!(fix_mode, flags::FixMode::Apply | flags::FixMode::Diff) {
            if let Ok(FixerResult {
                result,
                transformed,
                fixed,
            }) = lint_fix(
                path,
                package,
                noqa,
                unsafe_fixes,
                settings,
                &source_kind,
                source_type,
            ) {
                if !fixed.is_empty() {
                    match fix_mode {
                        flags::FixMode::Apply => transformed.write(&mut File::create(path)?)?,
                        flags::FixMode::Diff => {
                            write!(
                                &mut io::stdout().lock(),
                                "{}",
                                source_kind.diff(&transformed, Some(path)).unwrap()
                            )?;
                        }
                        flags::FixMode::Generate => {}
                    }
                }
                let transformed = if let Cow::Owned(transformed) = transformed {
                    transformed
                } else {
                    source_kind
                };
                (result, transformed, fixed)
            } else {
                // If we fail to fix, lint the original source code.
                let result = lint_only(
                    path,
                    package,
                    settings,
                    noqa,
                    &source_kind,
                    source_type,
                    ParseSource::None,
                );
                let transformed = source_kind;
                let fixed = FixTable::default();
                (result, transformed, fixed)
            }
        } else {
            let result = lint_only(
                path,
                package,
                settings,
                noqa,
                &source_kind,
                source_type,
                ParseSource::None,
            );
            let transformed = source_kind;
            let fixed = FixTable::default();
            (result, transformed, fixed)
        };

    let diagnostics = result.diagnostics;

    if let Some((cache, relative_path, key)) = caching {
        // `FixMode::Apply` and `FixMode::Diff` rely on side-effects (writing to disk,
        // and writing the diff to stdout, respectively). If a file has diagnostics
        // with fixes, we need to avoid reading from and writing to the cache in these
        // modes.
        let use_fixes = match fix_mode {
            flags::FixMode::Generate => true,
            flags::FixMode::Apply | flags::FixMode::Diff => fixed.is_empty(),
        };

        // We don't cache files with diagnostics.
        let linted = diagnostics.is_empty() && use_fixes;
        cache.set_linted(relative_path.to_owned(), &key, linted);
    }

    let notebook_indexes = if let SourceKind::IpyNotebook(notebook) = transformed {
        FxHashMap::from_iter([(path.to_string_lossy().to_string(), notebook.into_index())])
    } else {
        FxHashMap::default()
    };

    Ok(Diagnostics {
        inner: diagnostics,
        fixed: FixMap::from_iter([(fs::relativize_path(path), fixed)]),
        notebook_indexes,
    })
}

/// Generate `Diagnostic`s from source code content derived from stdin.
pub(crate) fn lint_stdin(
    path: Option<&Path>,
    package: Option<PackageRoot<'_>>,
    contents: String,
    settings: &Settings,
    noqa: flags::Noqa,
    fix_mode: flags::FixMode,
) -> Result<Diagnostics> {
    let (source_type, py_source_type) = match path
        .map(|path| settings.linter.extension.get_source_type(path))
        .unwrap_or_default()
    {
        SourceType::Toml(source_type @ (TomlSourceType::Pyproject | TomlSourceType::Ruff)) => {
            if !settings
                .linter
                .rules
                .iter_enabled()
                .any(|rule_code| rule_code.lint_source().is_toml())
            {
                return Ok(Diagnostics::default());
            }

            let path = path.unwrap();

            let (diagnostics, fixed) =
                if matches!(fix_mode, flags::FixMode::Apply | flags::FixMode::Diff) {
                    let TomlFixerResult {
                        diagnostics,
                        transformed,
                        fixed,
                    } = lint_fix_toml(
                        path,
                        &contents,
                        &settings.linter,
                        source_type,
                        settings.unsafe_fixes,
                    );

                    match fix_mode {
                        flags::FixMode::Apply => {
                            write!(&mut io::stdout().lock(), "{transformed}")?;
                        }
                        flags::FixMode::Diff => {
                            if !fixed.is_empty() {
                                write!(
                                    &mut io::stdout().lock(),
                                    "{}",
                                    SourceKindDiff::from_text(
                                        &contents,
                                        transformed.as_ref(),
                                        Some(path),
                                    )
                                )?;
                            }
                        }
                        flags::FixMode::Generate => {}
                    }

                    (diagnostics, fixed)
                } else {
                    (
                        lint_toml(path, &contents, &settings.linter, source_type),
                        FixTable::default(),
                    )
                };

            return Ok(Diagnostics {
                inner: diagnostics,
                fixed: FixMap::from_iter([(fs::relativize_path(path), fixed)]),
                notebook_indexes: FxHashMap::default(),
            });
        }

        SourceType::Toml(_) | SourceType::Markdown => return Ok(Diagnostics::default()),
        source_type @ SourceType::Python(py_source_type) => (source_type, py_source_type),
    };

    // Extract the sources from the file.
    let source_kind = match SourceKind::from_source_code(contents, source_type) {
        Ok(Some(source_kind)) => source_kind,
        Ok(None) => return Ok(Diagnostics::default()),
        Err(err) => {
            return Ok(Diagnostics::from_source_error(&err, path, &settings.linter));
        }
    };

    // Lint the inputs.
    let (LinterResult { diagnostics, .. }, transformed, fixed) =
        if matches!(fix_mode, flags::FixMode::Apply | flags::FixMode::Diff) {
            if let Ok(FixerResult {
                result,
                transformed,
                fixed,
            }) = lint_fix(
                path.unwrap_or_else(|| Path::new("-")),
                package,
                noqa,
                settings.unsafe_fixes,
                &settings.linter,
                &source_kind,
                py_source_type,
            ) {
                match fix_mode {
                    flags::FixMode::Apply => {
                        // Write the contents to stdout, regardless of whether any errors were fixed.
                        transformed.write(&mut io::stdout().lock())?;
                    }
                    flags::FixMode::Diff => {
                        // But only write a diff if it's non-empty.
                        if !fixed.is_empty() {
                            write!(
                                &mut io::stdout().lock(),
                                "{}",
                                source_kind.diff(&transformed, path).unwrap()
                            )?;
                        }
                    }
                    flags::FixMode::Generate => {}
                }
                let transformed = if let Cow::Owned(transformed) = transformed {
                    transformed
                } else {
                    source_kind
                };
                (result, transformed, fixed)
            } else {
                // If we fail to fix, lint the original source code.
                let result = lint_only(
                    path.unwrap_or_else(|| Path::new("-")),
                    package,
                    &settings.linter,
                    noqa,
                    &source_kind,
                    py_source_type,
                    ParseSource::None,
                );

                // Write the contents to stdout anyway.
                if fix_mode.is_apply() {
                    source_kind.write(&mut io::stdout().lock())?;
                }

                let transformed = source_kind;
                let fixed = FixTable::default();
                (result, transformed, fixed)
            }
        } else {
            let result = lint_only(
                path.unwrap_or_else(|| Path::new("-")),
                package,
                &settings.linter,
                noqa,
                &source_kind,
                py_source_type,
                ParseSource::None,
            );
            let transformed = source_kind;
            let fixed = FixTable::default();
            (result, transformed, fixed)
        };

    let notebook_indexes = if let SourceKind::IpyNotebook(notebook) = transformed {
        FxHashMap::from_iter([(
            path.map_or_else(|| "-".into(), |path| path.to_string_lossy().to_string()),
            notebook.into_index(),
        )])
    } else {
        FxHashMap::default()
    };

    Ok(Diagnostics {
        inner: diagnostics,
        fixed: FixMap::from_iter([(
            fs::relativize_path(path.unwrap_or_else(|| Path::new("-"))),
            fixed,
        )]),
        notebook_indexes,
    })
}
