use std::{ops::Range, sync::Arc};

use cloud_llm_client::{EditPredictionRejectReason, PredictEditsRequestTrigger};
use edit_prediction_types::{PredictedCursorPosition, interpolate_edits};
use gpui::{AsyncApp, Entity, SharedString};
use language::{Anchor, Buffer, BufferSnapshot, EditPreview, TextBufferSnapshot};
use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput};

#[derive(Clone, Default, Debug, PartialEq, Eq, Hash)]
pub struct EditPredictionId(pub SharedString);

impl From<EditPredictionId> for gpui::ElementId {
    fn from(value: EditPredictionId) -> Self {
        gpui::ElementId::Name(value.0)
    }
}

impl std::fmt::Display for EditPredictionId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

#[derive(Clone, serde::Serialize)]
#[serde(tag = "version", content = "input")]
pub enum EditPredictionInputs {
    V2(Zeta2PromptInput),
    V3(Zeta3PromptInput),
}

/// A prediction response that was returned from the provider, whether it was ultimately valid or not.
pub struct EditPredictionResult {
    pub prediction: EditPrediction,
    pub reject_reason: Option<EditPredictionRejectReason>,
    pub e2e_latency: std::time::Duration,
}

impl EditPredictionResult {
    pub async fn new(
        id: EditPredictionId,
        edited_buffer: &Entity<Buffer>,
        edited_buffer_snapshot: &BufferSnapshot,
        edits: Arc<[(Range<Anchor>, Arc<str>)]>,
        cursor_position: Option<PredictedCursorPosition>,
        editable_range: Option<Range<Anchor>>,
        inputs: EditPredictionInputs,
        model_version: Option<String>,
        trigger: PredictEditsRequestTrigger,
        e2e_latency: std::time::Duration,
        cx: &mut AsyncApp,
    ) -> Self {
        let (edits, reject_reason, new_snapshot): (
            Arc<[(Range<Anchor>, Arc<str>)]>,
            Option<EditPredictionRejectReason>,
            Option<BufferSnapshot>,
        ) = if edits.is_empty() {
            (
                Arc::default(),
                Some(EditPredictionRejectReason::Empty),
                None,
            )
        } else {
            edited_buffer.read_with(cx, |buffer, _cx| {
                let new_snapshot = buffer.snapshot();
                match interpolate_edits(&edited_buffer_snapshot, &new_snapshot, &edits) {
                    Some(edits) if edits.is_empty() => (
                        Arc::default(),
                        Some(EditPredictionRejectReason::InterpolatedEmpty),
                        None,
                    ),
                    Some(edits) => (Arc::from(edits), None, Some(new_snapshot)),
                    None => (
                        Arc::default(),
                        Some(EditPredictionRejectReason::InterpolateFailed),
                        None,
                    ),
                }
            })
        };
        let snapshot = new_snapshot.unwrap_or_else(|| edited_buffer_snapshot.clone());

        let edit_preview = if !edits.is_empty() {
            edited_buffer
                .read_with(cx, |buffer, cx| buffer.preview_edits(edits.clone(), cx))
                .await
        } else {
            EditPreview::unchanged(edited_buffer_snapshot)
        };

        Self {
            prediction: EditPrediction {
                id,
                edits,
                cursor_position,
                editable_range,
                snapshot,
                edit_preview,
                inputs,
                buffer: edited_buffer.clone(),
                model_version,
                trigger,
            },
            reject_reason,
            e2e_latency,
        }
    }

    pub fn new_rejected(
        id: EditPredictionId,
        edited_buffer: &Entity<Buffer>,
        edited_buffer_snapshot: &BufferSnapshot,
        inputs: EditPredictionInputs,
        model_version: Option<String>,
        trigger: PredictEditsRequestTrigger,
        e2e_latency: std::time::Duration,
        reject_reason: EditPredictionRejectReason,
    ) -> Self {
        Self {
            prediction: EditPrediction {
                id,
                edits: Arc::default(),
                cursor_position: None,
                editable_range: None,
                snapshot: edited_buffer_snapshot.clone(),
                edit_preview: EditPreview::unchanged(edited_buffer_snapshot),
                inputs,
                buffer: edited_buffer.clone(),
                model_version,
                trigger,
            },
            reject_reason: Some(reject_reason),
            e2e_latency,
        }
    }
}

#[derive(Clone)]
pub struct EditPrediction {
    pub id: EditPredictionId,
    pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
    pub cursor_position: Option<PredictedCursorPosition>,
    pub editable_range: Option<Range<Anchor>>,
    pub snapshot: BufferSnapshot,
    pub edit_preview: EditPreview,
    pub inputs: EditPredictionInputs,
    pub buffer: Entity<Buffer>,
    pub model_version: Option<String>,
    pub trigger: PredictEditsRequestTrigger,
}

impl EditPrediction {
    pub fn interpolate(
        &self,
        new_snapshot: &TextBufferSnapshot,
    ) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
        interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
    }

    pub fn targets_buffer(&self, buffer: &Buffer) -> bool {
        self.snapshot.remote_id() == buffer.remote_id()
    }
}

impl std::fmt::Debug for EditPrediction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EditPrediction")
            .field("id", &self.id)
            .field("edits", &self.edits)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use super::*;
    use gpui::{App, Entity, TestAppContext, prelude::*};
    use language::{Buffer, ToOffset as _};
    use zeta_prompt::Zeta2PromptInput;

    #[gpui::test]
    async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
        let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
        let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
            to_prediction_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
        });

        let edit_preview = cx
            .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
            .await;

        let prediction = EditPrediction {
            id: EditPredictionId("prediction-1".into()),
            edits,
            cursor_position: None,
            editable_range: None,
            snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
            buffer: buffer.clone(),
            edit_preview,
            model_version: None,
            trigger: PredictEditsRequestTrigger::Other,
            inputs: EditPredictionInputs::V2(Zeta2PromptInput {
                events: vec![],
                related_files: Some(vec![]),
                active_buffer_diagnostics: vec![],
                cursor_path: Path::new("path.txt").into(),
                cursor_offset_in_excerpt: 0,
                cursor_excerpt: "".into(),
                excerpt_start_row: None,
                excerpt_ranges: Default::default(),
                syntax_ranges: None,
                in_open_source_repo: false,
                can_collect_data: false,
                repo_url: None,
            }),
        };

        cx.update(|cx| {
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(2..5, "REM".into()), (9..11, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(2..2, "REM".into()), (6..8, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.undo(cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(2..5, "REM".into()), (9..11, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(3..3, "EM".into()), (7..9, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(4..4, "M".into()), (8..10, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(9..11, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(4..4, "M".into()), (8..10, "".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
            assert_eq!(
                from_prediction_edits(
                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                    &buffer,
                    cx
                ),
                vec![(4..4, "M".into())]
            );

            buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
            assert_eq!(prediction.interpolate(&buffer.read(cx).snapshot()), None);
        })
    }

    fn to_prediction_edits(
        iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
        buffer: &Entity<Buffer>,
        cx: &App,
    ) -> Vec<(Range<Anchor>, Arc<str>)> {
        let buffer = buffer.read(cx);
        iterator
            .into_iter()
            .map(|(range, text)| {
                (
                    buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
                    text,
                )
            })
            .collect()
    }

    fn from_prediction_edits(
        editor_edits: &[(Range<Anchor>, Arc<str>)],
        buffer: &Entity<Buffer>,
        cx: &App,
    ) -> Vec<(Range<usize>, Arc<str>)> {
        let buffer = buffer.read(cx);
        editor_edits
            .iter()
            .map(|(range, text)| {
                (
                    range.start.to_offset(buffer)..range.end.to_offset(buffer),
                    text.clone(),
                )
            })
            .collect()
    }
}