use std::borrow::Cow;

use lsp_types::FoldingRangeRequest;
use lsp_types::{FoldingRange, FoldingRangeKind, FoldingRangeParams, Uri};
use ruff_db::source::source_text;
use ruff_text_size::TextRange;
use ty_ide::folding_ranges;
use ty_project::{ProjectDatabase, SemanticDb as _};

use crate::db::Db;
use crate::document::ToRangeExt;
use crate::server::api::traits::{
    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
};
use crate::session::DocumentSnapshot;
use crate::session::client::Client;

pub(crate) struct FoldingRangeRequestHandler;

impl RequestHandler for FoldingRangeRequestHandler {
    type RequestType = FoldingRangeRequest;
}

impl BackgroundDocumentRequestHandler for FoldingRangeRequestHandler {
    fn document_uri(params: &FoldingRangeParams) -> Cow<'_, Uri> {
        Cow::Borrowed(&params.text_document.uri)
    }

    fn run_with_snapshot(
        db: &ProjectDatabase,
        snapshot: &DocumentSnapshot,
        _client: &Client,
        _params: FoldingRangeParams,
    ) -> crate::server::Result<Option<Vec<FoldingRange>>> {
        if snapshot
            .workspace_settings()
            .is_language_services_disabled()
        {
            return Ok(None);
        }

        let Some(file) = snapshot.to_notebook_or_file(db) else {
            return Ok(None);
        };

        // If this document is a notebook cell, only return folding ranges for this cell.
        // Folding ranges are reported as local ranges, so returning notebook-wide ranges
        // would produce invalid coordinates for every other cell.
        let mut cell_range: Option<TextRange> = None;

        if snapshot.document().is_cell()
            && let Some(notebook_document) = db.notebook_document(file)
            && let Some(notebook) = source_text(db, file).as_notebook()
        {
            let cell_index = notebook_document.cell_index_by_uri(snapshot.uri());
            cell_range = cell_index.and_then(|index| notebook.cell_range(index));
        }

        let results: Vec<_> = folding_ranges(db, db.program_file(file).python_file(db), cell_range)
            .into_iter()
            .filter_map(|folding_range| {
                let lsp_range = folding_range
                    .range
                    .to_lsp_range(db, file, snapshot.encoding())?;

                let kind = folding_range.kind.map(|k| match k {
                    ty_ide::FoldingRangeKind::Comment => FoldingRangeKind::Comment,
                    ty_ide::FoldingRangeKind::Imports => FoldingRangeKind::Imports,
                    ty_ide::FoldingRangeKind::Region => FoldingRangeKind::Region,
                });

                Some(FoldingRange {
                    start_line: lsp_range.local_range().start.line,
                    start_character: Some(lsp_range.local_range().start.character),
                    end_line: lsp_range.local_range().end.line,
                    end_character: Some(lsp_range.local_range().end.character),
                    kind,
                    collapsed_text: None,
                })
            })
            .collect();

        if results.is_empty() {
            Ok(None)
        } else {
            Ok(Some(results))
        }
    }
}

impl RetriableRequestHandler for FoldingRangeRequestHandler {}
