use crate::{AgentTool, ToolCallEventStream, ToolInput};
use agent_client_protocol as acp;
use gpui::{App, SharedString, Task};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::sync::Arc;

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
#[schemars(inline)]
pub enum PlanEntryStatus {
    /// The task has not started yet.
    Pending,
    /// The task is currently being worked on.
    InProgress,
    /// The task has been successfully completed.
    Completed,
}

impl From<PlanEntryStatus> for acp::PlanEntryStatus {
    fn from(value: PlanEntryStatus) -> Self {
        match value {
            PlanEntryStatus::Pending => acp::PlanEntryStatus::Pending,
            PlanEntryStatus::InProgress => acp::PlanEntryStatus::InProgress,
            PlanEntryStatus::Completed => acp::PlanEntryStatus::Completed,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
pub struct PlanItem {
    /// Human-readable description of what this task aims to accomplish.
    pub step: String,
    /// The current status of this task.
    pub status: PlanEntryStatus,
}

impl From<PlanItem> for acp::PlanEntry {
    fn from(value: PlanItem) -> Self {
        acp::PlanEntry::new(
            value.step,
            acp::PlanEntryPriority::Medium,
            value.status.into(),
        )
    }
}

/// Updates the task plan.
///
/// Provide a list of plan entries, each with a step and status.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
pub struct UpdatePlanToolInput {
    /// The list of plan entries and their current statuses.
    pub plan: Vec<PlanItem>,
}

pub struct UpdatePlanTool;

impl UpdatePlanTool {
    fn to_plan(input: UpdatePlanToolInput) -> acp::Plan {
        acp::Plan::new(input.plan.into_iter().map(Into::into).collect())
    }
}

impl AgentTool for UpdatePlanTool {
    type Input = UpdatePlanToolInput;
    type Output = String;

    const NAME: &'static str = "update_plan";

    fn kind() -> acp::ToolKind {
        acp::ToolKind::Think
    }

    fn initial_title(
        &self,
        input: Result<Self::Input, serde_json::Value>,
        _cx: &mut App,
    ) -> SharedString {
        match input {
            Ok(input) if input.plan.is_empty() => "Clear plan".into(),
            Ok(_) | Err(_) => "Update plan".into(),
        }
    }

    fn run(
        self: Arc<Self>,
        input: ToolInput<Self::Input>,
        event_stream: ToolCallEventStream,
        cx: &mut App,
    ) -> Task<Result<Self::Output, Self::Output>> {
        cx.spawn(async move |_cx| {
            let input = input
                .recv()
                .await
                .map_err(|e| format!("Failed to receive tool input: {e}"))?;

            event_stream.update_plan(Self::to_plan(input));

            Ok("Plan updated".to_string())
        })
    }

    fn replay(
        &self,
        input: Self::Input,
        _output: Self::Output,
        event_stream: ToolCallEventStream,
        _cx: &mut App,
    ) -> anyhow::Result<()> {
        event_stream.update_plan(Self::to_plan(input));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ToolCallEventStream;
    use gpui::TestAppContext;
    use pretty_assertions::assert_eq;

    fn sample_input() -> UpdatePlanToolInput {
        UpdatePlanToolInput {
            plan: vec![
                PlanItem {
                    step: "Inspect the existing tool wiring".to_string(),
                    status: PlanEntryStatus::Completed,
                },
                PlanItem {
                    step: "Implement the update_plan tool".to_string(),
                    status: PlanEntryStatus::InProgress,
                },
                PlanItem {
                    step: "Add tests".to_string(),
                    status: PlanEntryStatus::Pending,
                },
            ],
        }
    }

    #[gpui::test]
    async fn test_run_emits_plan_event(cx: &mut TestAppContext) {
        let tool = Arc::new(UpdatePlanTool);
        let (event_stream, mut event_rx) = ToolCallEventStream::test();

        let input = sample_input();
        let result = cx
            .update(|cx| tool.run(ToolInput::resolved(input.clone()), event_stream, cx))
            .await
            .expect("tool should succeed");

        assert_eq!(result, "Plan updated".to_string());

        let plan = event_rx.expect_plan().await;
        assert_eq!(
            plan,
            acp::Plan::new(vec![
                acp::PlanEntry::new(
                    "Inspect the existing tool wiring",
                    acp::PlanEntryPriority::Medium,
                    acp::PlanEntryStatus::Completed,
                ),
                acp::PlanEntry::new(
                    "Implement the update_plan tool",
                    acp::PlanEntryPriority::Medium,
                    acp::PlanEntryStatus::InProgress,
                ),
                acp::PlanEntry::new(
                    "Add tests",
                    acp::PlanEntryPriority::Medium,
                    acp::PlanEntryStatus::Pending,
                ),
            ])
        );
    }

    #[gpui::test]
    async fn test_replay_emits_plan_event(cx: &mut TestAppContext) {
        let tool = UpdatePlanTool;
        let (event_stream, mut event_rx) = ToolCallEventStream::test();

        let input = sample_input();

        cx.update(|cx| {
            tool.replay(input.clone(), "Plan updated".to_string(), event_stream, cx)
                .expect("replay should succeed");
        });

        let plan = event_rx.expect_plan().await;
        assert_eq!(
            plan,
            acp::Plan::new(vec![
                acp::PlanEntry::new(
                    "Inspect the existing tool wiring",
                    acp::PlanEntryPriority::Medium,
                    acp::PlanEntryStatus::Completed,
                ),
                acp::PlanEntry::new(
                    "Implement the update_plan tool",
                    acp::PlanEntryPriority::Medium,
                    acp::PlanEntryStatus::InProgress,
                ),
                acp::PlanEntry::new(
                    "Add tests",
                    acp::PlanEntryPriority::Medium,
                    acp::PlanEntryStatus::Pending,
                ),
            ])
        );
    }

    #[gpui::test]
    async fn test_initial_title(cx: &mut TestAppContext) {
        let tool = UpdatePlanTool;

        let title = cx.update(|cx| tool.initial_title(Ok(sample_input()), cx));
        assert_eq!(title, SharedString::from("Update plan"));

        let title =
            cx.update(|cx| tool.initial_title(Ok(UpdatePlanToolInput { plan: Vec::new() }), cx));
        assert_eq!(title, SharedString::from("Clear plan"));
    }
}