//! Per-session state and the turn loop.
//!
//! `run_turn` mirrors `hagent`'s headless turn orchestration (build request →
//! stream → fold events → snapshot), but emits ACP `SessionUpdate`s instead of
//! printing. It runs inside a `cx.spawn`ed task (never directly in the `prompt`
//! handler), so it may block on outbound requests (e.g. `request_permission`)
//! without deadlocking the dispatch loop.

use std::collections::HashMap;
use std::sync::{Arc, Mutex as StdMutex};

use agent_client_protocol as acp;
use agent_client_protocol::ConnectionTo;
use agent_client_protocol::schema::{
    ConfigOptionUpdate, ContentBlock, ContentChunk, CurrentModeUpdate, SessionId,
    SessionNotification, SessionUpdate, StopReason, TextContent, ToolCall, ToolCallContent,
    ToolCallId, ToolCallStatus, ToolKind,
};
use futures_util::StreamExt;
use hagent_core::{
    CallStatus, Client, ConvLogger, Conversation, RenderedBlock, Settings, SnapshotMeta, TokenStore,
};
use tokio::sync::Mutex as TokioMutex;
use tokio_util::sync::CancellationToken;

use crate::translate::{self, TranslateState};
use crate::{oauth, permission};

/// A live session: the immutable HTTP client (used to stream turns) plus the
/// mutable conversation/settings behind an async mutex, and the in-flight
/// cancellation token behind a sync mutex (so `session/cancel` can trigger it
/// without contending for the turn lock).
pub struct SessionHandle {
    /// The HTTP client is immutable for the session's life, so it lives outside
    /// the turn lock — letting the SSE stream borrow it while the turn loop
    /// mutably borrows the locked `SessionState`.
    pub client: Client,
    pub state: TokioMutex<SessionState>,
    pub cancel: StdMutex<Option<CancellationToken>>,
}

/// Mutable per-session state, guarded by [`SessionHandle::state`].
pub struct SessionState {
    pub conv: Conversation,
    pub settings: Settings,
    pub logger: Option<Arc<ConvLogger>>,
    pub created_ms: u64,
    /// Remembered MCP approval decisions (`allow_always`/`reject_always`) keyed
    /// by server label (Phase 3).
    pub remembered: HashMap<String, bool>,
    /// Model ids offered in the model config-option dropdown (fetched once from
    /// the server's `/v1/models`), so rebuilding the option after a change needs
    /// no extra round-trip.
    pub model_choices: Vec<String>,
    /// Whether the model picker groups by provider (vs. a flat list).
    pub group_models: bool,
}

impl SessionHandle {
    pub fn new(client: Client, state: SessionState) -> Self {
        Self {
            client,
            state: TokioMutex::new(state),
            cancel: StdMutex::new(None),
        }
    }

    /// Cancel the in-flight turn, if any. Called from the `session/cancel`
    /// notification handler.
    pub fn trigger_cancel(&self) {
        if let Some(tok) = self.cancel.lock().unwrap().as_ref() {
            tok.cancel();
        }
    }
}

/// Run one prompt turn to completion, streaming `SessionUpdate`s to the client.
/// Returns the `StopReason` the `prompt` handler reports back. All errors are
/// surfaced as message chunks + a stop reason — this never returns `Err`, so the
/// spawned task that drives it never tears down the connection.
pub async fn run_turn(
    session: &Arc<SessionHandle>,
    cx: &ConnectionTo<acp::Client>,
    store: &TokenStore,
    session_id: SessionId,
    user_text: Option<String>,
) -> StopReason {
    let mut guard = session.state.lock().await;
    let st = &mut *guard;
    let mut tool_state = TranslateState::default();

    // Slash commands (e.g. `/model …`, `/effort …`, `/status`) are handled
    // locally — apply the change, confirm, refresh the editor's pickers, and end
    // the turn without calling upstream.
    if let Some(text) = user_text.as_deref()
        && let Some(reply) = handle_slash_command(st, text)
    {
        emit_message(cx, &session_id, &reply);
        refresh_pickers(st, cx, &session_id);
        write_snapshot(st);
        return StopReason::EndTurn;
    }

    // MCP OAuth token gate, BEFORE the first request — Hadrian never persists the
    // bearer, so it's re-resolved every turn and folded into resolved_authorization.
    oauth::ensure_mcp_auth(&mut st.settings.mcp_servers, store, cx, &session_id).await;

    // First iteration carries the user's prompt (recorded to the transcript via
    // `push_user` and sent as an input item via `build_request`); continuation
    // turns (after approvals) carry `None`.
    let mut pending = user_text;

    let final_stop = loop {
        if let Some(text) = pending.as_ref() {
            st.conv.push_user(text.clone());
        }
        let req = st.conv.build_request(&st.settings, pending.take());

        let token = CancellationToken::new();
        *session.cancel.lock().unwrap() = Some(token.clone());

        let mut stream = match session.client.stream(req).await {
            Ok(s) => s,
            Err(e) => {
                let msg = e.to_string();
                st.conv.fail(msg.clone());
                emit_error(cx, &session_id, &msg);
                *session.cancel.lock().unwrap() = None;
                break StopReason::Refusal;
            }
        };

        let mut stop = StopReason::EndTurn;
        loop {
            tokio::select! {
                biased;
                () = token.cancelled() => { stop = StopReason::Cancelled; break; }
                next = stream.next() => match next {
                    Some(Ok(event)) => {
                        if let Some(update) = translate::to_update(&event, &mut tool_state) {
                            let _ = cx.send_notification(
                                SessionNotification::new(session_id.clone(), update),
                            );
                        }
                        translate::note_stop(&event, &mut stop);
                        st.conv.apply(event);
                    }
                    Some(Err(e)) => {
                        let msg = e.to_string();
                        st.conv.fail(msg.clone());
                        emit_error(cx, &session_id, &msg);
                        stop = StopReason::Refusal;
                        break;
                    }
                    None => break,
                }
            }
        }
        *session.cancel.lock().unwrap() = None;

        if matches!(stop, StopReason::Cancelled | StopReason::Refusal) {
            break stop;
        }

        // Resolve any MCP approvals the model requested, then run a continuation
        // turn: `build_request` drains the decided approvals into
        // `mcp_approval_response` input items.
        if st.conv.has_undecided_approvals() {
            resolve_approvals(st, cx, &session_id).await;
            pending = None;
            continue;
        }

        break stop;
    };

    // Reactive re-auth: an auth-shaped error on an `mcp_call` flags the server;
    // clear its token so the next turn re-prompts for login.
    for label in st.conv.take_mcp_auth_failures() {
        if let Some(server) = st
            .settings
            .mcp_servers
            .iter()
            .find(|s| s.server_label == label)
        {
            let _ = store.delete(&server.server_url);
        }
    }

    write_snapshot(st);
    final_stop
}

/// Ask the client to decide each undecided MCP approval (auto-resolving servers
/// the user chose to always allow/reject). The decisions are drained into the
/// continuation turn by `build_request`.
async fn resolve_approvals(
    st: &mut SessionState,
    cx: &ConnectionTo<acp::Client>,
    session_id: &SessionId,
) {
    while st.conv.has_undecided_approvals() {
        let Some(approval) = st.conv.current_approval().cloned() else {
            break;
        };
        let approve = match st.remembered.get(&approval.server_label).copied() {
            Some(decision) => decision,
            None => {
                permission::request_decision(cx, session_id, &approval, &mut st.remembered).await
            }
        };
        st.conv.decide_approval(approve);
    }
}

/// Persist the resumable snapshot after a settled turn (best-effort).
fn write_snapshot(st: &SessionState) {
    let Some(logger) = st.logger.clone() else {
        return;
    };
    let meta = SnapshotMeta {
        conversation_id: logger.conversation_id(),
        base_url: &st.settings.base_url,
        model: &st.settings.model,
        skill_id: st.settings.skill_id.as_deref(),
        created_ms: st.created_ms,
    };
    logger.write_snapshot(&st.conv.capture_snapshot(meta));
}

/// Replay a restored transcript to the client as `SessionUpdate`s so the editor
/// repaints prior history after `session/load`. Fire-and-forget (no awaits), so
/// it is safe to call directly from the load handler.
pub fn replay_transcript(
    cx: &ConnectionTo<acp::Client>,
    session_id: &SessionId,
    transcript: &[RenderedBlock],
) {
    for (i, block) in transcript.iter().enumerate() {
        for update in block_to_updates(i, block) {
            let _ = cx.send_notification(SessionNotification::new(session_id.clone(), update));
        }
    }
}

/// Map one restored transcript block to the `SessionUpdate`(s) that repaint it.
/// Tool calls replay as single completed `ToolCall`s (no follow-up updates), so
/// their ids are synthetic (`replay-<index>`).
fn block_to_updates(i: usize, block: &RenderedBlock) -> Vec<SessionUpdate> {
    match block {
        RenderedBlock::UserMsg(t) => vec![SessionUpdate::UserMessageChunk(ContentChunk::new(
            t.clone().into(),
        ))],
        RenderedBlock::AssistantText { text, .. } => vec![SessionUpdate::AgentMessageChunk(
            ContentChunk::new(text.clone().into()),
        )],
        RenderedBlock::Reasoning { text } => vec![SessionUpdate::AgentThoughtChunk(
            ContentChunk::new(text.clone().into()),
        )],
        RenderedBlock::ShellCall {
            command,
            output,
            status,
            ..
        } => {
            let mut content = vec![ToolCallContent::from(command.clone())];
            if !output.is_empty() {
                content.push(ToolCallContent::from(output.clone()));
            }
            let title = command.lines().next().unwrap_or("shell").to_string();
            vec![replay_tool(i, title, ToolKind::Execute, status, content)]
        }
        RenderedBlock::WebSearch { status } => {
            vec![replay_tool(
                i,
                "web search".into(),
                ToolKind::Fetch,
                status,
                vec![],
            )]
        }
        RenderedBlock::FileSearch { queries, status } => {
            let content = if queries.is_empty() {
                vec![]
            } else {
                vec![ToolCallContent::from(queries.join("\n"))]
            };
            vec![replay_tool(
                i,
                "file search".into(),
                ToolKind::Search,
                status,
                content,
            )]
        }
        RenderedBlock::McpCall {
            server_label,
            name,
            output,
            status,
            ..
        } => {
            let content = output
                .clone()
                .filter(|s| !s.is_empty())
                .map(|s| vec![ToolCallContent::from(s)])
                .unwrap_or_default();
            vec![replay_tool(
                i,
                format!("{server_label}::{name}"),
                ToolKind::Other,
                status,
                content,
            )]
        }
        RenderedBlock::Error { message, .. } => vec![message_chunk(&format!("[error] {message}"))],
        RenderedBlock::Notice(t) => vec![message_chunk(t)],
        // Approval prompts and compaction markers don't repaint as content.
        RenderedBlock::ApprovalPrompt { .. } | RenderedBlock::CompactionMarker => vec![],
    }
}

fn replay_tool(
    i: usize,
    title: String,
    kind: ToolKind,
    status: &CallStatus,
    content: Vec<ToolCallContent>,
) -> SessionUpdate {
    let mut tc = ToolCall::new(ToolCallId::new(format!("replay-{i}")), title)
        .kind(kind)
        .status(map_call_status(status));
    if !content.is_empty() {
        tc = tc.content(content);
    }
    SessionUpdate::ToolCall(tc)
}

fn map_call_status(status: &CallStatus) -> ToolCallStatus {
    match status {
        CallStatus::Completed => ToolCallStatus::Completed,
        CallStatus::Failed => ToolCallStatus::Failed,
        CallStatus::InProgress => ToolCallStatus::InProgress,
        CallStatus::Pending => ToolCallStatus::Pending,
    }
}

fn message_chunk(text: &str) -> SessionUpdate {
    SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(TextContent::new(
        text.to_string(),
    ))))
}

/// Surface a transport/stream error to the editor as an agent message chunk.
fn emit_error(cx: &ConnectionTo<acp::Client>, session_id: &SessionId, msg: &str) {
    emit_message(cx, session_id, &format!("[error] {msg}"));
}

/// Send a plain agent message chunk to the editor.
fn emit_message(cx: &ConnectionTo<acp::Client>, session_id: &SessionId, msg: &str) {
    let update = SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
        TextContent::new(format!("\n{msg}")),
    )));
    let _ = cx.send_notification(SessionNotification::new(session_id.clone(), update));
}

/// Handle a `/command` typed (or picked) in the composer. Returns the
/// confirmation text to show, or `None` if the text isn't a recognised command
/// (in which case it's sent upstream as a normal prompt). A leading `/` is
/// required so ordinary prompts beginning with "model"/"status" aren't hijacked.
fn handle_slash_command(st: &mut SessionState, text: &str) -> Option<String> {
    let body = text.trim().strip_prefix('/')?;
    let mut parts = body.splitn(2, char::is_whitespace);
    let cmd = parts.next()?;
    let arg = parts.next().unwrap_or("").trim();
    match cmd {
        "model" => {
            if arg.is_empty() {
                Some(format!("Current model: {}", st.settings.model))
            } else {
                st.settings.model = arg.to_string();
                Some(format!("Model set to `{arg}`."))
            }
        }
        "effort" => {
            if arg.is_empty() {
                Some(format!("Current reasoning effort: {}", current_effort(st)))
            } else {
                crate::agent::apply_effort(&mut st.settings, arg);
                Some(format!("Reasoning effort set to `{arg}`."))
            }
        }
        "status" => Some(status_report(st)),
        // Unknown `/command` — fall through to a normal prompt.
        _ => None,
    }
}

/// Emit fresh `ConfigOptionUpdate` + `CurrentModeUpdate` so the editor's pickers
/// reflect a settings change made via slash command.
fn refresh_pickers(st: &SessionState, cx: &ConnectionTo<acp::Client>, session_id: &SessionId) {
    let options = crate::agent::config_options(&st.settings, &st.model_choices, st.group_models);
    let _ = cx.send_notification(SessionNotification::new(
        session_id.clone(),
        SessionUpdate::ConfigOptionUpdate(ConfigOptionUpdate::new(options)),
    ));
    let _ = cx.send_notification(SessionNotification::new(
        session_id.clone(),
        SessionUpdate::CurrentModeUpdate(CurrentModeUpdate::new(current_effort(st))),
    ));
}

fn current_effort(st: &SessionState) -> String {
    st.settings
        .reasoning
        .as_ref()
        .and_then(|r| r.effort.clone())
        .unwrap_or_else(|| "off".to_string())
}

fn status_report(st: &SessionState) -> String {
    let tools: Vec<&str> = st
        .settings
        .tools
        .iter()
        .map(|t| match t {
            hagent_core::ToolKind::Shell => "shell",
            hagent_core::ToolKind::WebSearch => "web_search",
            hagent_core::ToolKind::FileSearch => "file_search",
            hagent_core::ToolKind::Mcp => "mcp",
        })
        .collect();
    let totals = &st.conv.totals;
    format!(
        "**model** {} · **effort** {} · **tools** [{}]\n**tokens** ↑{} ↓{} (total {})",
        st.settings.model,
        current_effort(st),
        tools.join(", "),
        totals.input_tokens,
        totals.output_tokens,
        totals.total_tokens,
    )
}
