use std::marker::PhantomData;
use std::sync::Arc;

use parking_lot::{Mutex, MutexGuard};
use rustc_hash::{FxBuildHasher, FxHashSet};
use salsa::{Durability, Setter};

use ruff_db::diagnostic::Diagnostic;
use ruff_db::files::File;

use crate::Project;
use crate::db::Db;

/// The indexed files of a project.
///
/// The indexing happens lazily, but the files are then cached for subsequent reads.
///
/// ## Implementation
/// The implementation uses internal mutability to transition between the lazy and indexed state
/// without triggering a new salsa revision. This is safe because the initial indexing happens on first access,
/// so no query can be depending on the contents of the indexed files before that. All subsequent mutations to
/// the indexed files must go through `IndexedMut`, which uses the Salsa setter `project.set_file_set` to
/// ensure that Salsa always knows when the set of indexed files have changed.
#[derive(Debug, get_size2::GetSize)]
pub struct IndexedFiles {
    // This mutex is intentionally non-poisoning: Salsa cancellation can unwind while the lazy
    // index is being built. Indexing only transitions to `Indexed` after a complete walk, so the
    // `Lazy` state remains valid and can be retried.
    state: Mutex<State>,
}

impl IndexedFiles {
    pub fn lazy() -> Self {
        Self {
            state: Mutex::new(State::Lazy),
        }
    }

    fn indexed(inner: Arc<IndexedInner>) -> Self {
        Self {
            state: Mutex::new(State::Indexed(inner)),
        }
    }

    pub(super) fn get(&self) -> Index<'_> {
        let state = self.state.lock();

        match &*state {
            State::Lazy => Index::Lazy(LazyFiles { files: state }),
            State::Indexed(inner) => Index::Indexed(Indexed {
                inner: Arc::clone(inner),
                _lifetime: PhantomData,
            }),
        }
    }

    pub(super) fn is_lazy(&self) -> bool {
        matches!(*self.state.lock(), State::Lazy)
    }

    /// Permanently freezes the project's file-set input without cloning the indexed files.
    pub(super) fn freeze(db: &mut dyn Db, project: Project) {
        let state = {
            let files = project.file_set(db);
            std::mem::replace(&mut *files.state.lock(), State::Lazy)
        };

        project
            .set_file_set(db)
            .with_durability(Durability::NEVER_CHANGE)
            .to(Self {
                state: Mutex::new(state),
            });
    }

    /// Returns a mutable view on the index that allows cheap in-place mutations.
    ///
    /// The changes are automatically written back to the database once the view is dropped.
    pub(super) fn indexed_mut(db: &mut dyn Db, project: Project) -> Option<IndexedMut<'_>> {
        // Calling `trigger_cancellation` cancels all pending salsa queries. This ensures that there are no pending
        // reads to the file set (this `db` is the only alive db).
        db.trigger_cancellation();

        // Replace the state with lazy. The `IndexedMut` guard restores the state
        // to `State::Indexed`  or sets a new `PackageFiles` when it gets dropped to ensure the state
        // is restored to how it has been before replacing the value.
        //
        // It isn't necessary to hold on to the lock after this point:
        // * The above call to `trigger_cancellation` guarantees that there's exactly **one** DB reference.
        // * `Indexed` has a `'db` lifetime, and this method requires a `&mut db`.
        //   This means that there can't be any pending reference to `Indexed` because Rust
        //   doesn't allow borrowing `db` as mutable (to call this method) and immutable (`Indexed<'db>`) at the same time.
        //   There can't be any other `Indexed<'db>` references created by clones of this DB because
        //   all clones must have been dropped at this point and the `Indexed`
        //   can't outlive the database (constrained by the `db` lifetime).
        let state = {
            let files = project.file_set(db);
            let mut locked = files.state.lock();
            std::mem::replace(&mut *locked, State::Lazy)
        };

        let indexed = match state {
            // If it's already lazy, just return. We also don't need to restore anything because the
            // replace above was a no-op.
            State::Lazy => return None,
            State::Indexed(indexed) => indexed,
        };

        Some(IndexedMut {
            db: Some(db),
            project,
            indexed,
            did_change: false,
        })
    }
}

impl Default for IndexedFiles {
    fn default() -> Self {
        Self::lazy()
    }
}

#[derive(Debug, get_size2::GetSize)]
enum State {
    /// The files of a package haven't been indexed yet.
    Lazy,

    /// The files are indexed. Stores the known files of a package.
    Indexed(Arc<IndexedInner>),
}

pub(super) enum Index<'db> {
    /// The index has not yet been computed. Allows inserting the files.
    Lazy(LazyFiles<'db>),

    Indexed(Indexed<'db>),
}

/// Package files that have not been indexed yet.
pub(super) struct LazyFiles<'db> {
    files: MutexGuard<'db, State>,
}

impl<'db> LazyFiles<'db> {
    /// Sets the indexed files of a package to `files`.
    pub(super) fn set(
        mut self,
        files: Vec<IndexedFile>,
        diagnostics: Vec<Diagnostic>,
    ) -> Indexed<'db> {
        let mut inner = IndexedInner {
            files: FxHashSet::with_capacity_and_hasher(files.len(), FxBuildHasher),
            scripts: FxHashSet::default(),
            diagnostics,
        };
        for IndexedFile { file, is_script } in files {
            inner.files.insert(file);
            if is_script {
                inner.scripts.insert(file);
            }
        }
        inner.files.shrink_to_fit();
        inner.scripts.shrink_to_fit();
        let files = Indexed {
            inner: Arc::new(inner),
            _lifetime: PhantomData,
        };
        *self.files = State::Indexed(Arc::clone(&files.inner));
        files
    }
}

/// A file classified by the walker without resolving its script settings or environment.
pub(crate) struct IndexedFile {
    pub(crate) file: File,
    pub(crate) is_script: bool,
}

/// The indexed files of the project.
///
/// Note: This type is intentionally non-cloneable. Making it cloneable requires
/// revisiting the locking behavior in [`IndexedFiles::indexed_mut`].
#[derive(Debug)]
pub struct Indexed<'db> {
    inner: Arc<IndexedInner>,
    // Preserve the lifetime of `PackageFiles`.
    _lifetime: PhantomData<&'db ()>,
}

#[derive(Debug, get_size2::GetSize)]
struct IndexedInner {
    files: FxHashSet<File>,
    scripts: FxHashSet<File>,
    diagnostics: Vec<Diagnostic>,
}

impl Indexed<'_> {
    pub fn iter(&self) -> impl Iterator<Item = File> {
        self.inner.files.iter().copied()
    }

    pub fn contains(&self, file: File) -> bool {
        self.inner.files.contains(&file)
    }

    pub fn is_empty(&self) -> bool {
        self.inner.files.is_empty()
    }

    pub(super) fn diagnostics(&self) -> &[Diagnostic] {
        &self.inner.diagnostics
    }

    pub(super) fn len(&self) -> usize {
        self.inner.files.len()
    }

    pub(super) fn scripts(&self) -> &FxHashSet<File> {
        &self.inner.scripts
    }
}

pub(super) type IndexedIter<'a> = std::iter::Copied<std::collections::hash_set::Iter<'a, File>>;

impl<'a> IntoIterator for &'a Indexed<'_> {
    type Item = File;
    type IntoIter = IndexedIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.inner.files.iter().copied()
    }
}

/// A Mutable view of a project's indexed files.
///
/// Allows in-place mutation of the files without deep cloning the hash set.
/// The changes are written back when the mutable view is dropped.
pub(super) struct IndexedMut<'db> {
    db: Option<&'db mut dyn Db>,
    project: Project,
    indexed: Arc<IndexedInner>,
    did_change: bool,
}

impl IndexedMut<'_> {
    pub(super) fn insert(&mut self, file: File, is_script: bool) {
        let inner = self.inner_mut();
        let file_added = inner.files.insert(file);
        let script_changed = if is_script {
            inner.scripts.insert(file)
        } else {
            inner.scripts.remove(&file)
        };
        self.did_change |= file_added || script_changed;
    }

    pub(super) fn remove(&mut self, file: File) -> bool {
        if self.inner_mut().files.remove(&file) {
            self.inner_mut().scripts.remove(&file);
            self.did_change = true;
            true
        } else {
            false
        }
    }

    pub(super) fn set_diagnostics(&mut self, diagnostics: Vec<Diagnostic>) {
        self.inner_mut().diagnostics = diagnostics;
    }

    fn inner_mut(&mut self) -> &mut IndexedInner {
        Arc::get_mut(&mut self.indexed)
            .expect("All references to `FilesSet` should have been dropped")
    }

    fn set_impl(&mut self) {
        let Some(db) = self.db.take() else {
            return;
        };

        let indexed = Arc::clone(&self.indexed);

        if self.did_change {
            // If there are changes, set the new file_set to trigger a salsa revision change.
            self.project
                .set_file_set(db)
                .to(IndexedFiles::indexed(indexed));
        } else {
            // The `indexed_mut` replaced the `state` with Lazy. Restore it back to the indexed state.
            *self.project.file_set(db).state.lock() = State::Indexed(indexed);
        }
    }
}

impl Drop for IndexedMut<'_> {
    fn drop(&mut self) {
        self.set_impl();
    }
}

#[cfg(test)]
mod tests {
    use std::assert_matches;
    use std::time::{Duration, Instant};

    use salsa::{Database, Durability, EventKind};

    use crate::ProjectMetadata;
    use crate::db::Db;
    use crate::db::testing::TestDb;
    use ruff_db::files::system_path_to_file;
    use ruff_db::system::{DbWithWritableSystem as _, SystemPathBuf};

    #[test]
    fn re_entrance() -> anyhow::Result<()> {
        let metadata = ProjectMetadata::new("test", SystemPathBuf::from("/test"));
        let mut db = TestDb::new(metadata);

        db.write_file("/test/test.py", "")?;

        let project = db.project();

        let file = system_path_to_file(&db, "/test/test.py")?;
        let files = project.files(&db);
        assert!(files.contains(file));

        // Calling files a second time should not dead-lock.
        // This can e.g. happen when `check_file` iterates over all files and
        // `should_check_file` queries the open files.
        let files_2 = project.files(&db);
        assert_eq!(
            files_2.iter().collect::<Vec<_>>(),
            files.iter().collect::<Vec<_>>()
        );

        Ok(())
    }

    #[test]
    fn cancelled_file_indexing_recovers_file_set() -> anyhow::Result<()> {
        let metadata = ProjectMetadata::new("test", SystemPathBuf::from("/test"));
        let mut db = TestDb::new(metadata);
        db.write_files((0..10_000).map(|index| (format!("/test/test_{index}.py"), "")))?;
        db.take_salsa_events();

        let project = db.project();
        let indexing_db = db.clone();
        let indexing = std::thread::spawn(move || project.files(&indexing_db).len());

        // Wait until the lazy index has started walking files before requesting cancellation.
        let deadline = Instant::now() + Duration::from_secs(5);
        while !db
            .take_salsa_events()
            .iter()
            .any(|event| matches!(event.kind, EventKind::WillCheckCancellation))
        {
            assert!(Instant::now() < deadline, "file indexing did not start");
            std::thread::yield_now();
        }

        db.synthetic_write(Durability::LOW);
        let cancelled = match indexing.join() {
            Ok(indexed_files) => {
                // Indexing can finish between the cancellation check and the write. That's valid;
                // there is no interrupted index to recover in that case.
                assert_eq!(indexed_files, 10_000);
                return Ok(());
            }
            Err(cancelled) => cancelled,
        };
        assert_matches!(
            cancelled.downcast_ref::<salsa::Cancelled>(),
            Some(salsa::Cancelled::PendingWrite),
            "file indexing did not propagate the salsa cancellation"
        );

        // The next access must retry the incomplete lazy index.
        assert_eq!(project.files(&db).len(), 10_000);

        Ok(())
    }
}
