use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use futures::StreamExt;
use tokio_util::sync::CancellationToken;

use crate::event::{AppEvent, EventBus};
use crate::llm::{FinishReason, LlmEvent, LlmEventStream, ProviderError};
use crate::permission::{PermissionService, Ruleset};
use crate::store::Store;
use crate::tool::{
    ContextReporter, MetadataSink, PermissionHandle, SubagentSpawner, Tool, ToolCtx, ToolError,
    ToolRegistry,
};
use crate::types::{Message, MessageId, Part, PartBody, PartId, SessionId, TokenUsage, ToolState};

use super::doomloop::DoomLoopGuard;
use super::jobs::JobBoard;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StepResult {
    Continue,
    Stop,
    Compact,
}

pub struct StepOutcome {
    pub result: StepResult,
    pub message_id: Option<MessageId>,
    pub usage: TokenUsage,
    /// Dollar cost of this step's usage (0.0 when no pricing is available).
    pub cost: f64,
    pub aborted: bool,
    /// Whether a file-mutating tool (`edit`/`write`) ran this step — drives the optional
    /// `after_file_tool` reminder injection on the next turn.
    pub used_file_tool: bool,
}

pub struct StepError {
    pub source: ProviderError,
    /// If `true`, the assistant message/parts were already persisted for this step —
    /// the outer loop must surface this as a terminal error, not retry.
    pub any_persisted: bool,
}

/// Everything a step needs that stays constant across the run (cheap to construct per step;
/// owns clones/handles, not the session itself, so the loop keeps ownership of `Session`).
pub struct StepContext {
    pub store: Store,
    pub bus: EventBus,
    pub tools: ToolRegistry,
    pub permissions: Arc<PermissionService>,
    pub static_rules: Ruleset,
    pub extra_rules: Arc<Mutex<Ruleset>>,
    /// Parent-effective ruleset for a subagent session; empty for a root session. Enables
    /// permission intersection on this session's tool calls.
    pub parent_rules: Ruleset,
    pub session_id: SessionId,
    pub cwd: PathBuf,
    /// Session's `tool-output` spill directory (see `tool::truncate`).
    pub data_dir: PathBuf,
    pub cancel: CancellationToken,
    /// Wall-clock for `created_at` stamps — passed in so tests stay deterministic.
    pub now: i64,
    /// Lets the `task` tool spawn subagents. `None` disables delegation (headless/tests).
    pub spawner: Option<Arc<dyn SubagentSpawner>>,
    /// This session's background job board (as a parent). Injected into requests when the
    /// running agent can delegate; `None` disables the board.
    pub job_board: Option<Arc<JobBoard>>,
    /// Present in subagent sessions: reports read files to this session's job on the parent
    /// board. `None` for root sessions (nothing to report to).
    pub context_reporter: Option<Arc<dyn ContextReporter>>,
    /// Language-server diagnostics source shared by the session's edit/write tool calls.
    pub diagnostics: Option<Arc<dyn crate::lsp::DiagnosticsSource>>,
}

struct FlushTracker {
    last_flush: Instant,
    bytes_since_flush: usize,
}

impl FlushTracker {
    fn new() -> Self {
        Self {
            last_flush: Instant::now(),
            bytes_since_flush: 0,
        }
    }

    /// docs/02-engine.md: flush to store every 50ms or 2KB of buffered delta text.
    fn should_flush(&mut self, new_bytes: usize) -> bool {
        self.bytes_since_flush += new_bytes;
        if self.bytes_since_flush >= 2048 || self.last_flush.elapsed() >= Duration::from_millis(50)
        {
            self.bytes_since_flush = 0;
            self.last_flush = Instant::now();
            true
        } else {
            false
        }
    }
}

/// Tool names that mutate files — after one runs, the optional `after_file_tool` reminder fires.
const FILE_TOOLS: &[&str] = &["edit", "write"];

struct Run<'a> {
    ctx: &'a StepContext,
    assistant: Option<Message>,
    used_file_tool: bool,
    next_idx: u32,
    active_text: Option<PartId>,
    active_reasoning: Option<PartId>,
    pending_tools: HashMap<String, PartId>,
    text_buf: HashMap<PartId, String>,
    reasoning_buf: HashMap<PartId, String>,
    tool_input_buf: HashMap<String, String>,
    flushers: HashMap<PartId, FlushTracker>,
    parts_by_id: HashMap<PartId, Part>,
}

impl<'a> Run<'a> {
    fn new(ctx: &'a StepContext) -> Self {
        Self {
            ctx,
            assistant: None,
            used_file_tool: false,
            next_idx: 0,
            active_text: None,
            active_reasoning: None,
            pending_tools: HashMap::new(),
            text_buf: HashMap::new(),
            reasoning_buf: HashMap::new(),
            tool_input_buf: HashMap::new(),
            flushers: HashMap::new(),
            parts_by_id: HashMap::new(),
        }
    }

    fn message_id(&self) -> MessageId {
        self.assistant
            .as_ref()
            .expect("assistant created")
            .id
            .clone()
    }

    async fn ensure_assistant(
        &mut self,
        model: crate::types::ModelRef,
        agent: &str,
    ) -> Result<(), ProviderError> {
        if self.assistant.is_some() {
            return Ok(());
        }
        let message =
            Message::new_assistant(self.ctx.session_id.clone(), model, agent, self.ctx.now);
        self.ctx
            .store
            .upsert_message(message.clone())
            .await
            .map_err(|e| ProviderError::Decode(e.to_string()))?;
        self.ctx.bus.publish(AppEvent::MessageCreated {
            message: message.clone(),
        });

        let step_start = Part {
            id: PartId::new(),
            message_id: message.id.clone(),
            session_id: self.ctx.session_id.clone(),
            idx: self.next_idx,
            body: PartBody::StepStart,
        };
        self.next_idx += 1;
        self.persist_part(step_start).await?;

        self.assistant = Some(message);
        Ok(())
    }

    async fn persist_part(&mut self, part: Part) -> Result<(), ProviderError> {
        self.ctx
            .store
            .upsert_part(part.clone())
            .await
            .map_err(|e| ProviderError::Decode(e.to_string()))?;
        self.ctx
            .bus
            .publish(AppEvent::PartUpdated { part: part.clone() });
        self.parts_by_id.insert(part.id.clone(), part);
        Ok(())
    }

    async fn on_text_start(&mut self, id: String) -> Result<(), ProviderError> {
        let part_id = PartId::from(id);
        let part = Part {
            id: part_id.clone(),
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: self.next_idx,
            body: PartBody::Text {
                text: String::new(),
                synthetic: false,
            },
        };
        self.next_idx += 1;
        self.text_buf.insert(part_id.clone(), String::new());
        self.flushers.insert(part_id.clone(), FlushTracker::new());
        self.active_text = Some(part_id);
        self.persist_part(part).await
    }

    async fn on_text_delta(&mut self, id: String, text: String) -> Result<(), ProviderError> {
        let part_id = PartId::from(id);
        let buf = self.text_buf.entry(part_id.clone()).or_default();
        buf.push_str(&text);
        let full_text = buf.clone();

        self.ctx.bus.publish(AppEvent::PartDelta {
            part_id: part_id.clone(),
            message_id: self.message_id(),
            delta: text.clone(),
        });

        let should_flush = self
            .flushers
            .entry(part_id.clone())
            .or_insert_with(FlushTracker::new)
            .should_flush(text.len());
        if should_flush {
            let part = Part {
                id: part_id,
                message_id: self.message_id(),
                session_id: self.ctx.session_id.clone(),
                idx: 0, // overwritten below from cached idx
                body: PartBody::Text {
                    text: full_text,
                    synthetic: false,
                },
            };
            self.persist_existing(part).await?;
        }
        Ok(())
    }

    async fn on_text_end(&mut self, id: String) -> Result<(), ProviderError> {
        let part_id = PartId::from(id);
        let full_text = self.text_buf.get(&part_id).cloned().unwrap_or_default();
        let part = Part {
            id: part_id,
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: 0,
            body: PartBody::Text {
                text: full_text,
                synthetic: false,
            },
        };
        self.active_text = None;
        self.persist_existing(part).await
    }

    async fn on_reasoning_start(&mut self, id: String) -> Result<(), ProviderError> {
        let part_id = PartId::from(id);
        let part = Part {
            id: part_id.clone(),
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: self.next_idx,
            body: PartBody::Reasoning {
                text: String::new(),
                signature: None,
            },
        };
        self.next_idx += 1;
        self.reasoning_buf.insert(part_id.clone(), String::new());
        self.flushers.insert(part_id.clone(), FlushTracker::new());
        self.active_reasoning = Some(part_id);
        self.persist_part(part).await
    }

    async fn on_reasoning_delta(&mut self, id: String, text: String) -> Result<(), ProviderError> {
        let part_id = PartId::from(id);
        let buf = self.reasoning_buf.entry(part_id.clone()).or_default();
        buf.push_str(&text);
        let full_text = buf.clone();

        self.ctx.bus.publish(AppEvent::PartDelta {
            part_id: part_id.clone(),
            message_id: self.message_id(),
            delta: text.clone(),
        });

        let should_flush = self
            .flushers
            .entry(part_id.clone())
            .or_insert_with(FlushTracker::new)
            .should_flush(text.len());
        if should_flush {
            let part = Part {
                id: part_id,
                message_id: self.message_id(),
                session_id: self.ctx.session_id.clone(),
                idx: 0,
                body: PartBody::Reasoning {
                    text: full_text,
                    signature: None,
                },
            };
            self.persist_existing(part).await?;
        }
        Ok(())
    }

    async fn on_reasoning_end(
        &mut self,
        id: String,
        signature: Option<String>,
    ) -> Result<(), ProviderError> {
        let part_id = PartId::from(id);
        let full_text = self
            .reasoning_buf
            .get(&part_id)
            .cloned()
            .unwrap_or_default();
        let part = Part {
            id: part_id,
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: 0,
            body: PartBody::Reasoning {
                text: full_text,
                signature,
            },
        };
        self.active_reasoning = None;
        self.persist_existing(part).await
    }

    /// Re-persists a part that already exists, preserving its original `idx`.
    async fn persist_existing(&mut self, mut part: Part) -> Result<(), ProviderError> {
        if let Some(existing) = self.parts_by_id.get(&part.id) {
            part.idx = existing.idx;
        }
        self.persist_part(part).await
    }

    async fn on_tool_input_start(
        &mut self,
        call_id: String,
        name: String,
    ) -> Result<(), ProviderError> {
        let part_id = PartId::new();
        self.pending_tools.insert(call_id.clone(), part_id.clone());
        self.tool_input_buf.insert(call_id.clone(), String::new());
        let part = Part {
            id: part_id,
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: self.next_idx,
            body: PartBody::Tool {
                call_id,
                name,
                state: ToolState::Pending {
                    partial_input: String::new(),
                },
            },
        };
        self.next_idx += 1;
        self.persist_part(part).await
    }

    async fn on_tool_input_delta(
        &mut self,
        call_id: String,
        json: String,
    ) -> Result<(), ProviderError> {
        let buf = self.tool_input_buf.entry(call_id.clone()).or_default();
        buf.push_str(&json);
        if let Some(part_id) = self.pending_tools.get(&call_id).cloned() {
            self.ctx.bus.publish(AppEvent::PartDelta {
                part_id,
                message_id: self.message_id(),
                delta: json,
            });
        }
        Ok(())
    }

    async fn on_tool_call(
        &mut self,
        call_id: String,
        name: String,
        input: serde_json::Value,
        doomloop: &mut DoomLoopGuard,
    ) -> Result<(), ProviderError> {
        if FILE_TOOLS.contains(&name.as_str()) {
            self.used_file_tool = true;
        }
        let part_id = self.pending_tools.remove(&call_id).unwrap_or_default();
        let running = Part {
            id: part_id.clone(),
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: self.next_idx_for(&part_id),
            body: PartBody::Tool {
                call_id: call_id.clone(),
                name: name.clone(),
                state: ToolState::Running {
                    input: input.clone(),
                    title: None,
                    metadata: serde_json::Value::Null,
                },
            },
        };
        self.persist_existing(running).await?;

        let mut doom_blocked = None;
        if doomloop.record_and_check(&name, &input) {
            let cancel = self.ctx.cancel.child_token();
            if let Err(e) = self
                .ctx
                .permissions
                .force_ask(
                    &self.ctx.session_id,
                    crate::permission::AskInput {
                        permission: "doom_loop".into(),
                        pattern: name.clone(),
                        always_pattern: name.clone(),
                        metadata: serde_json::json!({"tool": name, "input": input}),
                    },
                    &cancel,
                )
                .await
            {
                doom_blocked = Some(format!("model appears stuck repeating {name}: {e}"));
            }
        }

        let tool = self.ctx.tools.get(&name);
        let (title, output, metadata, is_error) = if let Some(reason) = doom_blocked {
            ("error".to_string(), reason, serde_json::Value::Null, true)
        } else {
            match tool {
                None => (
                    "error".to_string(),
                    format!("unknown tool: {name}"),
                    serde_json::Value::Null,
                    true,
                ),
                Some(tool) => self.run_tool(tool, call_id.clone(), input.clone()).await,
            }
        };

        // docs/05-tools.md: cap tool output at 30k chars regardless of which tool produced it.
        let output = crate::tool::truncate::truncate(&output, &self.ctx.data_dir, &call_id)
            .map(|t| t.text)
            .unwrap_or(output);

        let final_state = if is_error {
            ToolState::Error {
                input,
                error: output.clone(),
            }
        } else {
            ToolState::Completed {
                input,
                title: title.clone(),
                output: output.clone(),
                metadata,
                duration_ms: 0,
            }
        };
        let final_part = Part {
            id: part_id,
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: 0,
            body: PartBody::Tool {
                call_id,
                name,
                state: final_state,
            },
        };
        self.persist_existing(final_part).await
    }

    fn next_idx_for(&mut self, part_id: &PartId) -> u32 {
        if let Some(existing) = self.parts_by_id.get(part_id) {
            existing.idx
        } else {
            let idx = self.next_idx;
            self.next_idx += 1;
            idx
        }
    }

    async fn run_tool(
        &self,
        tool: Arc<dyn Tool>,
        call_id: String,
        input: serde_json::Value,
    ) -> (String, String, serde_json::Value, bool) {
        let call_cancel = self.ctx.cancel.child_token();
        let (metadata_sink, _metadata_rx) = MetadataSink::channel();
        let ask = PermissionHandle::new(
            self.ctx.permissions.clone(),
            self.ctx.session_id.clone(),
            self.ctx.static_rules.clone(),
            self.ctx.extra_rules.clone(),
            call_cancel.clone(),
        )
        .with_parent_rules(self.ctx.parent_rules.clone());
        let tool_ctx = ToolCtx {
            session_id: self.ctx.session_id.clone(),
            message_id: self.message_id(),
            call_id: call_id.clone(),
            cwd: self.ctx.cwd.clone(),
            data_dir: self.ctx.data_dir.clone(),
            cancel: call_cancel.clone(),
            ask,
            metadata: metadata_sink,
            spawner: self.ctx.spawner.clone(),
            context_reporter: self.ctx.context_reporter.clone(),
            diagnostics: self.ctx.diagnostics.clone(),
        };

        let result = tokio::select! {
            res = tool.execute(input, tool_ctx) => res,
            _ = call_cancel.cancelled() => Err(ToolError::Cancelled),
        };

        match result {
            Ok(output) => (output.title, output.output, output.metadata, false),
            Err(err) => (
                "error".to_string(),
                err.to_string(),
                serde_json::Value::Null,
                true,
            ),
        }
    }

    async fn on_finish(
        &mut self,
        reason: FinishReason,
        usage: TokenUsage,
        cost: f64,
    ) -> Result<StepResult, ProviderError> {
        let part = Part {
            id: PartId::new(),
            message_id: self.message_id(),
            session_id: self.ctx.session_id.clone(),
            idx: self.next_idx,
            body: PartBody::StepFinish {
                usage,
                cost,
                reason: reason.clone(),
            },
        };
        self.next_idx += 1;
        self.persist_part(part).await?;

        let result = match reason {
            FinishReason::ToolCalls => StepResult::Continue,
            _ => StepResult::Stop,
        };

        if let Some(message) = self.assistant.as_mut() {
            message.finished = Some(reason);
            message.usage = usage;
            self.ctx
                .store
                .upsert_message(message.clone())
                .await
                .map_err(|e| ProviderError::Decode(e.to_string()))?;
            self.ctx.bus.publish(AppEvent::MessageUpdated {
                message: message.clone(),
            });
        }
        Ok(result)
    }

    async fn mark_errored(&mut self, source: &ProviderError) {
        if let Some(message) = self.assistant.as_mut() {
            message.error = Some(crate::types::MessageError::Provider(source.to_string()));
            if self.ctx.store.upsert_message(message.clone()).await.is_ok() {
                self.ctx.bus.publish(AppEvent::MessageUpdated {
                    message: message.clone(),
                });
            }
        }
    }
}

/// Consumes one provider stream end-to-end: persists parts/messages as events arrive and
/// executes tool calls inline, in stream order (opencode's processor.ts semantics).
pub async fn process_step(
    mut stream: LlmEventStream,
    ctx: &StepContext,
    model: crate::types::ModelRef,
    agent: &str,
    cost: Option<crate::types::ModelCost>,
    doomloop: &mut DoomLoopGuard,
) -> Result<StepOutcome, StepError> {
    let mut run = Run::new(ctx);
    let mut usage = TokenUsage::default();
    let mut step_cost = 0.0;
    let mut result = StepResult::Stop;

    loop {
        let item = tokio::select! {
            item = stream.next() => item,
            _ = ctx.cancel.cancelled() => {
                run.mark_errored(&ProviderError::Cancelled).await;
                return Ok(StepOutcome {
                    result: StepResult::Stop,
                    message_id: run.assistant.as_ref().map(|m| m.id.clone()),
                    usage,
                    cost: step_cost,
                    aborted: true,
                    used_file_tool: run.used_file_tool,
                });
            }
        };
        let Some(item) = item else { break };

        let event = match item {
            Ok(event) => event,
            Err(source) => {
                let any_persisted = run.assistant.is_some();
                run.mark_errored(&source).await;
                return Err(StepError {
                    source,
                    any_persisted,
                });
            }
        };

        run.ensure_assistant(model.clone(), agent)
            .await
            .map_err(|source| StepError {
                source,
                any_persisted: false,
            })?;

        let outcome = match event {
            LlmEvent::TextStart { id } => run.on_text_start(id).await,
            LlmEvent::TextDelta { id, text } => run.on_text_delta(id, text).await,
            LlmEvent::TextEnd { id } => run.on_text_end(id).await,
            LlmEvent::ReasoningStart { id } => run.on_reasoning_start(id).await,
            LlmEvent::ReasoningDelta { id, text } => run.on_reasoning_delta(id, text).await,
            LlmEvent::ReasoningEnd { id, signature } => run.on_reasoning_end(id, signature).await,
            LlmEvent::ToolInputStart { call_id, name } => {
                run.on_tool_input_start(call_id, name).await
            }
            LlmEvent::ToolInputDelta { call_id, json } => {
                run.on_tool_input_delta(call_id, json).await
            }
            LlmEvent::ToolCall {
                call_id,
                name,
                input,
            } => run.on_tool_call(call_id, name, input, doomloop).await,
            LlmEvent::Finish {
                reason,
                usage: finish_usage,
            } => {
                usage = finish_usage;
                step_cost = cost.map(|c| c.cost_of(&finish_usage)).unwrap_or(0.0);
                match run.on_finish(reason, finish_usage, step_cost).await {
                    Ok(r) => {
                        result = r;
                        Ok(())
                    }
                    Err(e) => Err(e),
                }
            }
        };

        if let Err(source) = outcome {
            return Err(StepError {
                source,
                any_persisted: true,
            });
        }
    }

    Ok(StepOutcome {
        result,
        message_id: run.assistant.as_ref().map(|m| m.id.clone()),
        usage,
        cost: step_cost,
        aborted: false,
        used_file_tool: run.used_file_tool,
    })
}