use std::collections::HashMap;

use galaxy_agent_core::{
    AgentEvent, RuntimeActivity, RuntimeActivityStatus, RuntimeCapabilities, StopReason, Usage,
};
use uuid::Uuid;
use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};

use crate::ai::agent::runtime_activity;
use crate::ai::bedrock::response_translator::{
    build_add_agent_output_message, build_append_text, build_create_task, build_stream_init,
    build_user_query_message,
};
use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage};

pub(crate) struct RuntimeResponseConfig {
    pub(crate) task_id: String,
    pub(crate) conversation_id: String,
    pub(crate) needs_create_task: bool,
    pub(crate) user_query: Option<String>,
    pub(crate) model_id: String,
    pub(crate) max_context_tokens: Option<u32>,
    pub(crate) capabilities: RuntimeCapabilities,
    pub(crate) empty_output_message: Option<String>,
}

/// Converts the provider-neutral runtime lifecycle into Galaxy's existing
/// transcript protocol. Runtime implementations never need to know about the
/// protobuf messages consumed by the UI.
pub(crate) struct RuntimeResponseTranslator {
    config: RuntimeResponseConfig,
    request_id: String,
    initialized: bool,
    text_message_id: Option<String>,
    reasoning_message_id: Option<String>,
    activity_message_ids: HashMap<String, String>,
    activities: HashMap<String, RuntimeActivity>,
    has_visible_output: bool,
    usage: Usage,
    context_usage: Option<(u64, u64)>,
}

impl RuntimeResponseTranslator {
    pub(crate) fn new(config: RuntimeResponseConfig) -> Self {
        Self {
            config,
            request_id: Uuid::new_v4().to_string(),
            initialized: false,
            text_message_id: None,
            reasoning_message_id: None,
            activity_message_ids: HashMap::new(),
            activities: HashMap::new(),
            has_visible_output: false,
            usage: Usage::default(),
            context_usage: None,
        }
    }

    pub(crate) fn translate(&mut self, event: AgentEvent) -> Result<Vec<ResponseEvent>, String> {
        let mut events = Vec::new();
        match event {
            AgentEvent::TurnStarted { .. } => self.initialize(&mut events),
            AgentEvent::TextDelta { text } => {
                self.initialize(&mut events);
                self.add_or_append_text(&text, &mut events);
            }
            AgentEvent::ReasoningDelta { text } => {
                self.initialize(&mut events);
                self.add_or_append_reasoning(&text, &mut events);
            }
            AgentEvent::ReasoningCompleted { text, .. } => {
                self.initialize(&mut events);
                if self.reasoning_message_id.is_none() && !text.is_empty() {
                    self.add_or_append_reasoning(&text, &mut events);
                }
            }
            AgentEvent::RuntimeActivityUpdated { activity } => {
                if self.config.capabilities.host_tool_execution {
                    return Err(
                        "a host-tool runtime emitted runtime-owned tool activity".to_owned()
                    );
                }
                self.initialize(&mut events);
                self.text_message_id = None;
                self.upsert_runtime_activity(activity, &mut events)?;
            }
            AgentEvent::ContextUsageUpdated {
                used_tokens,
                context_size,
            } => {
                if self.config.capabilities.host_managed_history {
                    return Err(
                        "a host-history runtime reported session context occupancy".to_owned()
                    );
                }
                self.context_usage = Some((used_tokens, context_size));
            }
            AgentEvent::UserInputAccepted { text } => {
                if !self.config.capabilities.steering {
                    return Err("a non-steerable runtime accepted live user input".to_owned());
                }
                self.initialize(&mut events);
                events.push(build_user_query_message(&self.config.task_id, &text));
                self.text_message_id = None;
                self.reasoning_message_id = None;
            }
            AgentEvent::RuntimeNotice { message } => {
                self.initialize(&mut events);
                self.text_message_id = None;
                self.add_or_append_text(&message, &mut events);
                self.text_message_id = None;
            }
            AgentEvent::UsageUpdated { usage } => {
                if !self.config.capabilities.host_managed_history {
                    return Err("a session runtime reported provider request usage".to_owned());
                }
                self.usage = usage;
            }
            AgentEvent::TurnStopped { reason } => {
                self.initialize(&mut events);
                if !self.has_visible_output && reason != StopReason::Cancelled {
                    if let Some(message) = self.config.empty_output_message.clone() {
                        self.add_or_append_text(&message, &mut events);
                    }
                }
                events.push(self.finished(reason));
            }
            AgentEvent::Tool { .. } => {
                let owner = if self.config.capabilities.host_tool_execution {
                    "host-tool runtime emitted an unhandled tool lifecycle event"
                } else {
                    "session runtime handed runtime-owned tool execution to Galaxy"
                };
                return Err(owner.to_owned());
            }
        }
        Ok(events)
    }

    pub(crate) fn startup_error(&mut self, message: &str) -> Vec<ResponseEvent> {
        let mut events = Vec::new();
        self.initialize(&mut events);
        self.text_message_id = None;
        self.add_or_append_text(message, &mut events);
        events.push(self.finished(StopReason::Refusal));
        events
    }

    pub(crate) fn begin_followup_turn(&mut self) {
        self.text_message_id = None;
        self.reasoning_message_id = None;
    }

    fn initialize(&mut self, events: &mut Vec<ResponseEvent>) {
        if self.initialized {
            return;
        }
        events.push(build_stream_init(
            &self.request_id,
            &self.config.conversation_id,
        ));
        if self.config.needs_create_task {
            events.push(build_create_task(&self.config.task_id));
        }
        if let Some(user_query) = &self.config.user_query {
            events.push(build_user_query_message(&self.config.task_id, user_query));
        }
        self.initialized = true;
    }

    fn add_or_append_text(&mut self, text: &str, events: &mut Vec<ResponseEvent>) {
        if text.is_empty() {
            return;
        }
        self.has_visible_output = true;
        if let Some(message_id) = &self.text_message_id {
            events.push(build_append_text(&self.config.task_id, message_id, text));
        } else {
            let message_id = Uuid::new_v4().to_string();
            events.push(build_add_agent_output_message(
                &self.config.task_id,
                &message_id,
                text,
            ));
            self.text_message_id = Some(message_id);
        }
    }

    fn add_or_append_reasoning(&mut self, text: &str, events: &mut Vec<ResponseEvent>) {
        if text.is_empty() {
            return;
        }
        self.has_visible_output = true;
        if let Some(message_id) = &self.reasoning_message_id {
            events.push(build_reasoning_message(
                &self.config.task_id,
                message_id,
                text,
                true,
            ));
        } else {
            let message_id = Uuid::new_v4().to_string();
            events.push(build_reasoning_message(
                &self.config.task_id,
                &message_id,
                text,
                false,
            ));
            self.reasoning_message_id = Some(message_id);
        }
    }

    fn upsert_runtime_activity(
        &mut self,
        activity: RuntimeActivity,
        events: &mut Vec<ResponseEvent>,
    ) -> Result<(), String> {
        let activity_id = activity.id.clone();
        let merged_activity = self
            .activities
            .entry(activity_id.clone())
            .or_insert_with(|| activity.clone());
        if !activity.title.trim().is_empty() {
            merged_activity.title = activity.title;
        }
        if activity.status.is_some() {
            merged_activity.status = activity.status;
        }
        if activity.output.is_some() {
            merged_activity.output = activity.output;
        }

        let server_message_data = runtime_activity::encode(merged_activity)
            .map_err(|error| format!("failed to encode runtime activity: {error}"))?;
        let fallback_text = runtime_activity_fallback_text(merged_activity);
        if let Some(message_id) = self.activity_message_ids.get(&activity_id) {
            events.push(build_update_runtime_activity_message(
                &self.config.task_id,
                message_id,
                &fallback_text,
                &server_message_data,
            ));
        } else {
            let message_id = Uuid::new_v4().to_string();
            events.push(build_add_runtime_activity_message(
                &self.config.task_id,
                &message_id,
                &fallback_text,
                &server_message_data,
            ));
            self.activity_message_ids.insert(activity_id, message_id);
        }
        self.has_visible_output = true;
        Ok(())
    }

    fn finished(&self, reason: StopReason) -> ResponseEvent {
        let reason = map_stop_reason(reason);
        if !self.config.capabilities.host_managed_history {
            let (used_tokens, context_size) = self.context_usage.unwrap_or_default();
            return build_context_finished(
                reason,
                &self.config.model_id,
                used_tokens,
                context_size,
            );
        }
        build_stream_finished(
            reason,
            StreamUsage {
                input_tokens: saturating_i32(self.usage.input_tokens),
                output_tokens: saturating_i32(self.usage.output_tokens),
                cache_read_tokens: saturating_i32(self.usage.cached_input_tokens),
                cache_write_tokens: saturating_i32(self.usage.cache_creation_input_tokens),
                cost_in_cents: 0.0,
                model_id: self.config.model_id.clone(),
                max_context_tokens: self.config.max_context_tokens,
            },
        )
    }
}

fn build_reasoning_message(
    task_id: &str,
    message_id: &str,
    text: &str,
    append: bool,
) -> ResponseEvent {
    let message = api::Message {
        id: message_id.to_owned(),
        task_id: task_id.to_owned(),
        request_id: String::new(),
        timestamp: None,
        server_message_data: String::new(),
        citations: Vec::new(),
        fetched_memories: Vec::new(),
        message: Some(api::message::Message::AgentReasoning(
            api::message::AgentReasoning {
                reasoning: text.to_owned(),
                finished_duration: None,
            },
        )),
    };
    let action = if append {
        api::client_action::Action::AppendToMessageContent(
            api::client_action::AppendToMessageContent {
                task_id: task_id.to_owned(),
                message: Some(message),
                mask: Some(prost_types::FieldMask {
                    paths: vec!["agent_reasoning.reasoning".to_owned()],
                }),
            },
        )
    } else {
        api::client_action::Action::AddMessagesToTask(api::client_action::AddMessagesToTask {
            task_id: task_id.to_owned(),
            messages: vec![message],
        })
    };
    runtime_client_action(action)
}

fn runtime_activity_fallback_text(activity: &RuntimeActivity) -> String {
    let title = &activity.title;
    let status = activity.status.as_ref().map(|status| match status {
        RuntimeActivityStatus::Pending => "waiting",
        RuntimeActivityStatus::InProgress => "running",
        RuntimeActivityStatus::Completed => "completed",
        RuntimeActivityStatus::Failed => "failed",
        RuntimeActivityStatus::Other(_) => "updated",
    });
    let mut text = match status {
        Some(status) => format!("> **{title}** — {status}"),
        None => format!("> **{title}**"),
    };
    if let Some(output) = &activity.output {
        text.push_str("\n\n");
        for line in output.lines() {
            text.push_str("    ");
            text.push_str(line);
            text.push('\n');
        }
    }
    text
}

fn build_add_runtime_activity_message(
    task_id: &str,
    message_id: &str,
    fallback_text: &str,
    server_message_data: &str,
) -> ResponseEvent {
    let message = runtime_activity_message(task_id, message_id, fallback_text, server_message_data);
    runtime_client_action(api::client_action::Action::AddMessagesToTask(
        api::client_action::AddMessagesToTask {
            task_id: task_id.to_owned(),
            messages: vec![message],
        },
    ))
}

fn build_update_runtime_activity_message(
    task_id: &str,
    message_id: &str,
    fallback_text: &str,
    server_message_data: &str,
) -> ResponseEvent {
    let message = runtime_activity_message(task_id, message_id, fallback_text, server_message_data);
    runtime_client_action(api::client_action::Action::UpdateTaskMessage(
        api::client_action::UpdateTaskMessage {
            task_id: task_id.to_owned(),
            message: Some(message),
            mask: Some(prost_types::FieldMask {
                paths: vec![
                    "agent_output.text".to_owned(),
                    "server_message_data".to_owned(),
                ],
            }),
        },
    ))
}

fn runtime_activity_message(
    task_id: &str,
    message_id: &str,
    fallback_text: &str,
    server_message_data: &str,
) -> api::Message {
    api::Message {
        id: message_id.to_owned(),
        task_id: task_id.to_owned(),
        request_id: String::new(),
        timestamp: None,
        server_message_data: server_message_data.to_owned(),
        citations: Vec::new(),
        fetched_memories: Vec::new(),
        message: Some(api::message::Message::AgentOutput(
            api::message::AgentOutput {
                text: fallback_text.to_owned(),
            },
        )),
    }
}

fn runtime_client_action(action: api::client_action::Action) -> ResponseEvent {
    ResponseEvent {
        r#type: Some(api::response_event::Type::ClientActions(
            api::response_event::ClientActions {
                actions: vec![ClientAction {
                    action: Some(action),
                }],
            },
        )),
    }
}

fn build_context_finished(
    reason: stream_finished::Reason,
    model_id: &str,
    used_tokens: u64,
    context_size: u64,
) -> ResponseEvent {
    let total_input_tokens = u32::try_from(used_tokens).unwrap_or(u32::MAX);
    let context_window_usage = if context_size == 0 {
        0.0
    } else {
        (used_tokens as f32 / context_size as f32).clamp(0.0, 1.0)
    };
    #[allow(deprecated)]
    let usage_metadata = stream_finished::ConversationUsageMetadata {
        context_window_usage,
        summarized: false,
        credits_spent: 0.0,
        platform_credits_spent: 0.0,
        total_input_tokens,
        token_usage: Vec::new(),
        tool_usage_metadata: None,
        warp_token_usage: HashMap::new(),
        byok_token_usage: HashMap::new(),
        custom_endpoint_token_usage: HashMap::new(),
        context_window_segments: Vec::new(),
    };
    ResponseEvent {
        r#type: Some(api::response_event::Type::Finished(
            api::response_event::StreamFinished {
                reason: Some(reason),
                token_usage: vec![stream_finished::TokenUsage {
                    model_id: model_id.to_owned(),
                    // Session runtimes report current occupancy, not tokens
                    // consumed by this individual request.
                    total_input: 0,
                    output: 0,
                    input_cache_read: 0,
                    input_cache_write: 0,
                    cost_in_cents: 0.0,
                }],
                should_refresh_model_config: false,
                request_cost: None,
                conversation_usage_metadata: Some(usage_metadata),
            },
        )),
    }
}

pub(crate) fn map_stop_reason(reason: StopReason) -> stream_finished::Reason {
    match reason {
        StopReason::Completed | StopReason::Cancelled => {
            stream_finished::Reason::Done(stream_finished::Done {})
        }
        StopReason::MaxTokens => {
            stream_finished::Reason::MaxTokenLimit(stream_finished::ReachedMaxTokenLimit {})
        }
        StopReason::ContextWindowExceeded => stream_finished::Reason::ContextWindowExceeded(
            stream_finished::ContextWindowExceeded {},
        ),
        StopReason::Refusal | StopReason::ToolLoopLimit | StopReason::Other(_) => {
            stream_finished::Reason::Other(stream_finished::Other {})
        }
    }
}

pub(crate) fn saturating_i32(value: u64) -> i32 {
    i32::try_from(value).unwrap_or(i32::MAX)
}

#[cfg(test)]
#[path = "event_translator_tests.rs"]
mod tests;