#![allow(
    clippy::expect_used,
    clippy::unwrap_used,
    clippy::panic,
    clippy::uninlined_format_args,
    clippy::collapsible_if,
    clippy::redundant_clone,
    clippy::needless_raw_string_hashes,
    clippy::single_match,
    clippy::redundant_closure_for_method_calls,
    clippy::redundant_pattern_matching,
    clippy::ignored_unit_patterns,
    clippy::clone_on_copy,
    clippy::manual_assert,
    clippy::unwrap_in_result,
    clippy::useless_vec
)]
//! Phase 0b binary — JSON-RPC gateway bridging SDK clients to the unified runtime.

use std::collections::HashMap;
use std::io::Write;
use std::sync::Arc;
use std::time::Duration;

use meerkat_mobkit::{
    AuthPolicy, Base64BlobStoreAdapter, BigQueryNaming, BinaryBlobStore, ConsolePolicy,
    DiscoverySpec, MOBKIT_CONTRACT_VERSION, MobBootstrapOptions, MobBootstrapSpec, MobKitConfig,
    ModuleConfig, ObjectStoreBlobStore, ReleaseMetadata, RestartPolicy, RuntimeDecisionState,
    RuntimeOpsPolicy, TrustedOidcRuntimeConfig, UnifiedRuntime, handle_mobkit_rpc_json,
    handle_unified_rpc_json, mob_handle_runtime::mob_definition_may_use_image_generation,
    start_mobkit_runtime,
};
use sha2::{Digest, Sha256};

use async_trait::async_trait;
use meerkat::{
    AgentEvent, AgentFactory, Config, CreateSessionRequest, EphemeralSessionService, FactoryAgent,
    FactoryAgentBuilder, SessionAgentBuilder, SessionError,
};
use meerkat_core::AgentToolDispatcher;
use meerkat_core::ContentBlock;
use meerkat_core::error::{AgentError, ToolError};
use meerkat_core::ops::ToolDispatchOutcome;
use meerkat_core::types::{ToolCallView, ToolDef, ToolResult};
use meerkat_mob::{MobDefinition, MobStorage};
use serde_json::{Value, json};
use tokio::sync::{Mutex, mpsc, oneshot};

fn minimal_decision_state() -> RuntimeDecisionState {
    RuntimeDecisionState {
        bigquery: BigQueryNaming {
            dataset: "default_dataset".to_string(),
            table: "default_table".to_string(),
        },
        modules: vec![],
        auth: AuthPolicy::default(),
        trusted_oidc: TrustedOidcRuntimeConfig {
            discovery_json: r#"{"issuer":"https://noop.example.com","authorization_endpoint":"https://noop.example.com/auth","token_endpoint":"https://noop.example.com/token","jwks_uri":"https://noop.example.com/.well-known/jwks.json","response_types_supported":["code"],"subject_types_supported":["public"],"id_token_signing_alg_values_supported":["RS256"]}"#.to_string(),
            jwks_json: r#"{"keys":[]}"#.to_string(),
            audience: "persistent-gateway".to_string(),
        },
        console: ConsolePolicy::default(),
        ops: RuntimeOpsPolicy::default(),
        release_metadata: ReleaseMetadata {
            targets: vec![
                "crates.io".to_string(),
                "npm".to_string(),
                "pypi".to_string(),
                "github-releases".to_string(),
            ],
            support_matrix: "lts".to_string(),
        },
    }
}

fn shell_module(id: &str, script: &str) -> ModuleConfig {
    ModuleConfig {
        id: id.to_string(),
        command: "sh".to_string(),
        args: vec!["-c".to_string(), script.to_string()],
        restart_policy: RestartPolicy::Never,
    }
}

/// Original single-shot mode: reads request from env, runs once, prints response.
fn run_single_shot() {
    let request = std::env::var("MOBKIT_RPC_REQUEST")
        .expect("MOBKIT_RPC_REQUEST must be set for rpc_gateway");

    let config = MobKitConfig {
        modules: vec![shell_module(
            "routing",
            r#"printf '%s\n' '{"event_id":"evt-routing","source":"module","timestamp_ms":101,"event":{"kind":"module","module":"routing","event_type":"ready","payload":{"family":"routing","health":{"state":"healthy"},"tools":{"list_method":"routing/tools.list","representative_call":{"method":"routing/tool.call","params_schema":{"tool":"string","input":"json"}}}}}}'"#,
        )],
        discovery: DiscoverySpec {
            namespace: "mobkit-rpc".to_string(),
            modules: vec!["routing".to_string()],
        },
        pre_spawn: vec![],
    };

    let mut runtime =
        start_mobkit_runtime(config, vec![], Duration::from_secs(1)).expect("runtime starts");
    let response = handle_mobkit_rpc_json(&mut runtime, &request, Duration::from_secs(1));
    print!("{response}");
    let _ = runtime.shutdown();
}

// ---------------------------------------------------------------------------
// StdioCallbackAgentBuilder — wraps FactoryAgentBuilder, sends callback/build_agent
// to Python over stdout before building the agent.
// ---------------------------------------------------------------------------

/// Shared handle for sending lines to stdout and receiving callback responses.
#[derive(Clone)]
struct StdioCallbackBridge {
    /// Send a line to stdout (the stdout writer task reads from this).
    stdout_tx: mpsc::Sender<String>,
    /// Pending callback responses keyed by callback ID.
    pending: Arc<Mutex<HashMap<String, oneshot::Sender<Value>>>>,
    /// Counter for generating unique callback IDs.
    counter: Arc<std::sync::atomic::AtomicU64>,
}

impl StdioCallbackBridge {
    fn new(stdout_tx: mpsc::Sender<String>) -> Self {
        Self {
            stdout_tx,
            pending: Arc::new(Mutex::new(HashMap::new())),
            counter: Arc::new(std::sync::atomic::AtomicU64::new(1)),
        }
    }

    /// Send a fire-and-forget notification to Python (no response expected).
    /// Uses `try_send` — may drop under backpressure.
    fn notify(&self, method: &str, params: Value) {
        let notification = json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        });
        if let Ok(line) = serde_json::to_string(&notification) {
            let _ = self.stdout_tx.try_send(line);
        }
    }

    /// Send a notification with reliable delivery (async, waits for channel space).
    /// Use for callbacks where delivery must not be silently lost.
    async fn notify_reliable(&self, method: &str, params: Value) {
        let notification = json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        });
        if let Ok(line) = serde_json::to_string(&notification) {
            if let Err(e) = self.stdout_tx.send(line).await {
                eprintln!("[mobkit-gateway] failed to deliver {method}: {e}");
            }
        }
    }

    /// Send a callback request to Python and wait for the response.
    async fn call(&self, method: &str, params: Value) -> Result<Value, String> {
        let id = self
            .counter
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let id_str = format!("cb-{id}");

        let (tx, rx) = oneshot::channel();
        self.pending.lock().await.insert(id_str.clone(), tx);

        let request = json!({
            "jsonrpc": "2.0",
            "id": id_str,
            "method": method,
            "params": params,
        });
        let line = match serde_json::to_string(&request) {
            Ok(l) => l,
            Err(e) => {
                self.pending.lock().await.remove(&id_str);
                return Err(e.to_string());
            }
        };
        if let Err(_) = self.stdout_tx.send(line).await {
            self.pending.lock().await.remove(&id_str);
            return Err("stdout channel closed".to_string());
        }

        // Wait for Python to respond (routed by the stdin multiplexer)
        match tokio::time::timeout(Duration::from_mins(2), rx).await {
            Ok(Ok(value)) => {
                if let Some(error) = value.get("error") {
                    Err(format!(
                        "callback error: {}",
                        error
                            .get("message")
                            .and_then(|m| m.as_str())
                            .unwrap_or("unknown")
                    ))
                } else {
                    Ok(value.get("result").cloned().unwrap_or(Value::Null))
                }
            }
            Ok(Err(_)) => Err("callback response channel dropped".to_string()),
            Err(_) => {
                self.pending.lock().await.remove(&id_str);
                Err("callback timed out after 120s".to_string())
            }
        }
    }

    /// Route an incoming callback response (has "id" starting with "cb-").
    async fn route_callback_response(&self, msg: Value) {
        let id = msg
            .get("id")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();
        if let Some(tx) = self.pending.lock().await.remove(&id) {
            let _ = tx.send(msg);
        }
    }
}

#[async_trait]
impl meerkat_mobkit::identity_first::gateway_bridges::CallbackBridge for StdioCallbackBridge {
    async fn call(&self, method: &str, params: Value) -> Result<Value, String> {
        self.call(method, params).await
    }
}

/// Tool dispatcher that routes tool calls to Python via the callback bridge.
///
/// Created from tool name strings provided by `SessionBuildOptions.add_tools()`.
/// When the agent calls a tool, `dispatch()` sends `callback/call_tool` to Python
/// and returns the result.
struct CallbackToolDispatcher {
    bridge: StdioCallbackBridge,
    scope_id: String,
    tool_defs: Arc<[Arc<ToolDef>]>,
}

impl CallbackToolDispatcher {
    fn new(bridge: StdioCallbackBridge, scope_id: String, tool_names: Vec<String>) -> Self {
        let tool_defs: Vec<Arc<ToolDef>> = tool_names
            .into_iter()
            .map(|name| {
                Arc::new(ToolDef {
                    name: name.into(),
                    description: "Python callback tool".to_string(),
                    input_schema: json!({"type": "object"}),
                    provenance: None,
                })
            })
            .collect();
        Self {
            bridge,
            scope_id,
            tool_defs: tool_defs.into(),
        }
    }
}

#[async_trait]
impl AgentToolDispatcher for CallbackToolDispatcher {
    fn tools(&self) -> Arc<[Arc<ToolDef>]> {
        Arc::clone(&self.tool_defs)
    }

    async fn dispatch(&self, call: ToolCallView<'_>) -> Result<ToolDispatchOutcome, ToolError> {
        let args: Value =
            serde_json::from_str(call.args.get()).map_err(|e| ToolError::InvalidArguments {
                name: call.name.to_string(),
                reason: e.to_string(),
            })?;
        let params = json!({
            "scope_id": self.scope_id,
            "tool": call.name,
            "arguments": args,
        });
        match self.bridge.call("callback/call_tool", params).await {
            Ok(result) => {
                let text = result
                    .get("content")
                    .map(|v| {
                        if let Some(s) = v.as_str() {
                            s.to_string()
                        } else {
                            serde_json::to_string(v).unwrap_or_default()
                        }
                    })
                    .unwrap_or_else(|| serde_json::to_string(&result).unwrap_or_default());
                Ok(ToolResult {
                    tool_use_id: call.id.to_string(),
                    content: vec![ContentBlock::Text { text }],
                    is_error: false,
                }
                .into())
            }
            Err(err) => Ok(ToolResult {
                tool_use_id: call.id.to_string(),
                content: vec![ContentBlock::Text {
                    text: format!("Tool execution failed: {err}"),
                }],
                is_error: true,
            }
            .into()),
        }
    }
}

/// Wraps FactoryAgentBuilder — sends callback/build_agent to Python before building.
struct StdioCallbackAgentBuilder {
    inner: FactoryAgentBuilder,
    bridge: StdioCallbackBridge,
    has_session_builder: bool,
    /// Session store for loading sessions by ID when the Python builder
    /// sets `resume_session_id`. Only populated in persistent mode.
    session_store: Option<Arc<dyn meerkat::SessionStore>>,
}

#[async_trait]
impl SessionAgentBuilder for StdioCallbackAgentBuilder {
    type Agent = FactoryAgent;

    async fn build_agent(
        &self,
        req: &CreateSessionRequest,
        event_tx: mpsc::Sender<AgentEvent>,
    ) -> Result<Self::Agent, SessionError> {
        if !self.has_session_builder {
            return self.inner.build_agent(req, event_tx).await;
        }

        // Generate a unique scope ID for this build — used to isolate tool
        // handlers between concurrent sessions on the Python side.
        let scope_id = format!(
            "build-{}",
            self.bridge
                .counter
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
        );

        // Send callback to Python with full session context.
        // app_context flows from SpawnMemberSpec.context → build.app_context.
        let options = json!({
            "scope_id": scope_id,
            "session_id": req.labels.as_ref().and_then(|l| l.get("session_id")),
            "model": &req.model,
            "prompt": &req.prompt,
            "labels": &req.labels,
            "app_context": req.build.as_ref()
                .and_then(|b| b.app_context.as_ref()),
        });
        let params = json!({ "options": options });
        let callback_result = self.bridge.call("callback/build_agent", params).await;

        match callback_result {
            Ok(result) => {
                // Apply Python-returned options to a cloned request
                let mut modified_req = CreateSessionRequest {
                    model: req.model.clone(),
                    prompt: req.prompt.clone(),
                    render_metadata: req.render_metadata.clone(),
                    system_prompt: req.system_prompt.clone(),
                    max_tokens: req.max_tokens,
                    event_tx: req.event_tx.clone(),
                    skill_references: req.skill_references.clone(),
                    initial_turn: req.initial_turn.clone(),
                    build: req.build.clone(),
                    labels: req.labels.clone(),
                    deferred_prompt_policy: req.deferred_prompt_policy,
                };
                // Apply additional_instructions as system prompt extension
                if let Some(instructions) = result.get("additional_instructions") {
                    if let Some(arr) = instructions.as_array() {
                        let combined: Vec<&str> = arr.iter().filter_map(|v| v.as_str()).collect();
                        if !combined.is_empty() {
                            let extra = combined.join("\n");
                            modified_req.system_prompt = Some(match &modified_req.system_prompt {
                                Some(existing) => format!("{existing}\n{extra}"),
                                None => extra,
                            });
                        }
                    }
                }
                // Apply labels
                if let Some(labels) = result.get("labels").and_then(|v| v.as_object()) {
                    let label_map = modified_req.labels.get_or_insert_with(Default::default);
                    for (k, v) in labels {
                        if let Some(s) = v.as_str() {
                            label_map.insert(k.clone(), s.to_string());
                        }
                    }
                }
                // Apply resume_session_id: Python builder can request resuming
                // an existing session. The gateway loads the session from the
                // store and sets it on build.resume_session.
                if let Some(resume_id) = result.get("resume_session_id").and_then(|v| v.as_str()) {
                    if let Some(ref store) = self.session_store {
                        let sid =
                            meerkat_core::types::SessionId::parse(resume_id).map_err(|_| {
                                SessionError::Agent(AgentError::ToolError(format!(
                                    "callback/build_agent: invalid resume_session_id: {resume_id}"
                                )))
                            })?;
                        // Validate against any spawn-level resume already set.
                        if let Some(existing) = modified_req
                            .build
                            .as_ref()
                            .and_then(|b| b.resume_session.as_ref())
                        {
                            if existing.id() != &sid {
                                return Err(SessionError::Agent(AgentError::ToolError(format!(
                                    "callback/build_agent: resume_session_id conflict: \
                                     spawn set {} but hook set {resume_id}",
                                    existing.id()
                                ))));
                            }
                            // Same ID — already loaded, skip.
                        } else {
                            let session = store.load(&sid).await.map_err(|e| {
                                SessionError::Agent(AgentError::ToolError(format!(
                                    "callback/build_agent: failed to load resume session {resume_id}: {e}"
                                )))
                            })?;
                            let session = session.ok_or_else(|| {
                                SessionError::Agent(AgentError::ToolError(format!(
                                    "callback/build_agent: resume session not found: {resume_id}"
                                )))
                            })?;
                            let build = modified_req.build.get_or_insert_with(|| {
                                meerkat_core::service::SessionBuildOptions::default()
                            });
                            build.resume_session = Some(session);
                        }
                    } else {
                        return Err(SessionError::Agent(AgentError::ToolError(
                            "callback/build_agent: resume_session_id requires persistent mode \
                             (no session store available in ephemeral mode)"
                                .to_string(),
                        )));
                    }
                }
                // Callback tools: Python SDK provides tool names via add_tools()
                // or register_tool(). Create a CallbackToolDispatcher that routes
                // tool calls back to Python via callback/call_tool.
                if let Some(tools) = result.get("tools") {
                    match tools.as_array() {
                        Some(arr) => {
                            let mut tool_names = Vec::with_capacity(arr.len());
                            for v in arr {
                                if let Some(name) = v.as_str() {
                                    tool_names.push(name.to_string());
                                } else {
                                    return Err(SessionError::Agent(AgentError::ToolError(
                                        format!(
                                            "callback/build_agent: tools must be strings, got: {v}"
                                        ),
                                    )));
                                }
                            }
                            if !tool_names.is_empty() {
                                let dispatcher = CallbackToolDispatcher::new(
                                    self.bridge.clone(),
                                    scope_id.clone(),
                                    tool_names,
                                );
                                let build = modified_req.build.get_or_insert_with(|| {
                                    meerkat_core::service::SessionBuildOptions::default()
                                });
                                build.external_tools = Some(Arc::new(dispatcher));
                            }
                        }
                        None => {
                            return Err(SessionError::Agent(AgentError::ToolError(format!(
                                "callback/build_agent: tools must be a JSON array, got: {tools}"
                            ))));
                        }
                    }
                }
                self.inner.build_agent(&modified_req, event_tx).await
            }
            Err(err) => {
                // Propagate the error — before_create failure aborts session creation.
                // This is an intentional breaking change from v0.5.x where failures
                // were silently swallowed with a fallback to default build.
                Err(SessionError::Agent(AgentError::ToolError(format!(
                    "callback/build_agent failed: {err}"
                ))))
            }
        }
    }
}

/// Persistent mode: reads JSON-RPC over stdin, bootstraps unified runtime, serves HTTP.
#[tokio::main]
async fn run_persistent() {
    use tokio::io::{AsyncBufReadExt, BufReader};

    // Initialize tracing subscriber so meerkat-mob/meerkat-runtime errors
    // are visible on stderr. Without this, all tracing events are silently
    // dropped and runtime failures (agent build, LLM calls, comms drain)
    // are invisible.
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
        )
        .with_writer(std::io::stderr)
        .with_ansi(false)
        .init();

    let stdin = tokio::io::stdin();
    let mut reader = BufReader::new(stdin);

    // 1. Read first line — must be mobkit/init
    let mut init_line = String::new();
    if reader.read_line(&mut init_line).await.unwrap_or(0) == 0 {
        eprintln!("rpc_gateway: stdin closed before init request");
        std::process::exit(1);
    }

    let init_raw: Value = match serde_json::from_str(init_line.trim()) {
        Ok(v) => v,
        Err(e) => {
            let error_response = json!({
                "jsonrpc": "2.0",
                "id": null,
                "error": { "code": -32700, "message": format!("Parse error: {e}") }
            });
            println!(
                "{}",
                serde_json::to_string(&error_response)
                    .unwrap_or_else(|_| r#"{"error":"serialization failed"}"#.to_string())
            );
            std::process::exit(1);
        }
    };

    let request_id = init_raw.get("id").cloned().unwrap_or(Value::Null);
    let method = init_raw
        .get("method")
        .and_then(|m| m.as_str())
        .unwrap_or("");
    if method != "mobkit/init" {
        let error_response = json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "error": { "code": -32600, "message": format!("Expected mobkit/init, got {method}") }
        });
        println!(
            "{}",
            serde_json::to_string(&error_response)
                .unwrap_or_else(|_| r#"{"error":"serialization failed"}"#.to_string())
        );
        std::process::exit(1);
    }

    let params = init_raw.get("params").cloned().unwrap_or_else(|| json!({}));

    // 2. Parse init params
    let mob_config_param = params.get("mob_config").and_then(|v| v.as_str());
    let is_workspace_config = mob_config_param.is_some();
    let mob_config_toml = mob_config_param.unwrap_or(
        r#"
[mob]
id = "persistent-gateway"

[profiles.default]
model = "gpt-5.2"
external_addressable = true
"#,
    );

    let definition = MobDefinition::from_toml(mob_config_toml).unwrap_or_else(|e| {
        let error_response = json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "error": { "code": -32602, "message": format!("Invalid mob_config TOML: {e}") }
        });
        println!(
            "{}",
            serde_json::to_string(&error_response)
                .unwrap_or_else(|_| r#"{"error":"serialization failed"}"#.to_string())
        );
        std::process::exit(1);
    });
    let image_generation = mob_definition_may_use_image_generation(&definition);

    // Validate profile model names against the catalog.
    // A wrong model name (e.g., "claude-sonnet-4-5-20250514" instead of "claude-sonnet-4-5")
    // silently fails at LLM call time with no observable error. Catch it here.
    {
        let catalog = meerkat_models::catalog::catalog();
        let known_models: std::collections::HashSet<&str> =
            catalog.iter().map(|entry| entry.id).collect();

        for (profile_name, binding) in &definition.profiles {
            let Some(profile) = binding.as_inline() else {
                continue; // Realm refs are resolved at runtime, not validated here.
            };
            if !known_models.contains(profile.model.as_str()) {
                let model = &profile.model;
                // Find similar model names for the error hint
                let prefix = model.split('-').take(3).collect::<Vec<_>>().join("-");
                let mut suggestions: Vec<&str> = known_models
                    .iter()
                    .filter(|m| {
                        m.starts_with(&prefix)
                            || model
                                .starts_with(&m.split('-').take(3).collect::<Vec<_>>().join("-"))
                    })
                    .copied()
                    .collect();
                suggestions.sort_unstable();
                suggestions.truncate(5);
                let hint = if suggestions.is_empty() {
                    String::new()
                } else {
                    format!(". Did you mean one of: {}?", suggestions.join(", "))
                };
                fail_init(
                    &request_id,
                    -32602,
                    format!("Profile '{profile_name}' uses unknown model '{model}'{hint}"),
                );
            }
        }
    }

    let modules: Vec<ModuleConfig> = params
        .get("modules")
        .and_then(|v| serde_json::from_value(v.clone()).ok())
        .unwrap_or_default();

    let discovery_modules: Vec<String> = modules.iter().map(|m| m.id.clone()).collect();
    let module_config = MobKitConfig {
        modules,
        discovery: DiscoverySpec {
            namespace: "persistent-gateway".to_string(),
            modules: discovery_modules,
        },
        pre_spawn: vec![],
    };

    let has_session_builder = params
        .get("has_session_builder")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    let persistent_state = params
        .get("persistent_state")
        .and_then(|v| v.as_str())
        .map(std::path::PathBuf::from);

    let has_roster_provider = params
        .get("has_roster_provider")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    let has_topology_provider = params
        .get("has_topology_provider")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    let has_agent_customizer = params
        .get("has_agent_customizer")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    // Warn about SDK-forwarded runtime_options that the gateway does not yet
    // consume.  This surfaces "accepted but ignored" config as visible noise
    // so integrators notice early rather than debugging silent no-ops.
    if let Some(runtime_options) = params.get("runtime_options").and_then(|v| v.as_object()) {
        let unrecognized: Vec<&String> = runtime_options.keys().collect();
        if !unrecognized.is_empty() {
            eprintln!(
                "[mobkit-gateway] warning: runtime_options keys [{}] were sent by the SDK but are not consumed by the gateway — this config has no effect",
                unrecognized
                    .iter()
                    .map(|k| k.as_str())
                    .collect::<Vec<_>>()
                    .join(", ")
            );
        }
    }

    // 3. Set up stdout writer channel for multiplexed output
    let (stdout_tx, mut stdout_rx) = mpsc::channel::<String>(64);
    let stdout_writer = tokio::spawn(async move {
        while let Some(line) = stdout_rx.recv().await {
            let mut stdout = std::io::stdout().lock();
            let _ = writeln!(stdout, "{line}");
            let _ = stdout.flush();
            drop(stdout); // release lock before next await
        }
    });

    // 4. Build callback bridge and start stdin multiplexer BEFORE bootstrap.
    // This ensures callback responses (e.g. callback/build_agent during discovery
    // spawn) are routed even while UnifiedRuntime::bootstrap is running.
    let bridge = StdioCallbackBridge::new(stdout_tx.clone());
    let (rpc_tx, mut rpc_rx) = mpsc::channel::<String>(64);

    let stdin_reader = tokio::spawn({
        let bridge = bridge.clone();
        let rpc_tx = rpc_tx.clone();
        async move {
            let mut line = String::new();
            loop {
                line.clear();
                match reader.read_line(&mut line).await {
                    Ok(0) => break, // EOF
                    Ok(_) => {}
                    Err(_) => break,
                }
                let trimmed = line.trim();
                if trimmed.is_empty() {
                    continue;
                }
                let msg: Value = match serde_json::from_str(trimmed) {
                    Ok(v) => v,
                    Err(_) => continue,
                };
                // Callback responses: "id" starts with "cb-" and no "method"
                let is_callback_response = msg
                    .get("id")
                    .and_then(|v| v.as_str())
                    .is_some_and(|id| id.starts_with("cb-"))
                    && msg.get("method").is_none();

                if is_callback_response {
                    bridge.route_callback_response(msg).await;
                } else {
                    // Queue RPC request for the dispatch loop
                    let _ = rpc_tx.send(trimmed.to_string()).await;
                }
            }
        }
    });

    /// Helper: send a JSON-RPC error response for the init request and exit.
    fn fail_init(request_id: &Value, code: i32, message: String) -> ! {
        let error_response = json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "error": { "code": code, "message": message }
        });
        let mut stdout = std::io::stdout().lock();
        let _ = writeln!(
            stdout,
            "{}",
            serde_json::to_string(&error_response)
                .unwrap_or_else(|_| r#"{"error":"serialization failed"}"#.to_string())
        );
        let _ = stdout.flush();
        std::process::exit(1);
    }

    // 5. Build session service with callback bridge.
    let (mob_spec, _temp_dir) = if let Some(ref state_path) = persistent_state {
        if let Err(e) = std::fs::create_dir_all(state_path) {
            fail_init(
                &request_id,
                -32603,
                format!("failed to create persistent state directory: {e}"),
            );
        }
        let sqlite_path = state_path.join("sessions.db");
        let session_store: Arc<dyn meerkat::SessionStore> =
            match meerkat_store::SqliteSessionStore::open(sqlite_path) {
                Ok(s) => Arc::new(s),
                Err(e) => fail_init(
                    &request_id,
                    -32603,
                    format!("failed to open SQLite session store: {e}"),
                ),
            };
        let mob_storage = match MobStorage::persistent(state_path.join("mob.redb")) {
            Ok(s) => s,
            Err(e) => fail_init(
                &request_id,
                -32603,
                format!("failed to open redb mob storage: {e}"),
            ),
        };
        let binary_blob_store: Arc<dyn BinaryBlobStore> =
            match ObjectStoreBlobStore::local(state_path.join("blobs")) {
                Ok(store) => Arc::new(store),
                Err(e) => fail_init(
                    &request_id,
                    -32603,
                    format!("failed to open binary blob store: {e}"),
                ),
            };
        let blob_store: Arc<dyn meerkat_core::BlobStore> =
            Arc::new(Base64BlobStoreAdapter::new(binary_blob_store.clone()));
        // Persistent runtime store at <state_path>/runtime.sqlite — must
        // be Some() on the session service so archive/retire can mutate
        // the authoritative session, and must be persistent so resume
        // works across gateway restart.
        let runtime_db_path = state_path.join("runtime.sqlite");
        let runtime_store: Arc<dyn meerkat_runtime::RuntimeStore> =
            match meerkat_runtime::store::SqliteRuntimeStore::new(&runtime_db_path) {
                Ok(store) => Arc::new(store),
                Err(err) => {
                    tracing::warn!(
                        path = %runtime_db_path.display(),
                        error = %err,
                        "failed to open SqliteRuntimeStore; falling back to InMemoryRuntimeStore. \
                         Sessions will not survive process restart and archive operations may fail.",
                    );
                    Arc::new(meerkat_runtime::InMemoryRuntimeStore::new())
                }
            };
        let adapter = Arc::new(meerkat_runtime::MeerkatMachine::persistent(
            Arc::clone(&runtime_store),
            Arc::clone(&blob_store),
        ));
        // Match the ephemeral path's capability mask — only comms is enabled
        // by default. Apps control additional capabilities via their mob
        // definition profiles, not the gateway factory.
        let mut factory = AgentFactory::new(state_path).builtins(false).comms(true);
        if image_generation {
            factory = factory.with_image_generation_machine(adapter.clone());
        }
        let mut inner_builder = FactoryAgentBuilder::new(factory, Config::default());
        inner_builder.default_session_store = Some(Arc::new(meerkat_store::StoreAdapter::new(
            session_store.clone(),
        )));
        inner_builder.default_blob_store = Some(blob_store.clone());
        let callback_builder = StdioCallbackAgentBuilder {
            inner: inner_builder,
            bridge: bridge.clone(),
            has_session_builder,
            session_store: Some(session_store.clone()),
        };
        let session_service: Arc<dyn meerkat_mob::MobSessionService> =
            Arc::new(meerkat_session::PersistentSessionService::new(
                callback_builder,
                16,
                session_store,
                Some(Arc::clone(&runtime_store)),
                blob_store,
            ));
        let mut spec = MobBootstrapSpec::new(definition, mob_storage, session_service)
            .with_options(MobBootstrapOptions {
                allow_ephemeral_sessions: true,
                notify_orchestrator_on_resume: true,
                default_llm_client: None,
            });
        spec.runtime_adapter = Some(adapter);
        spec.binary_blob_store = Some(binary_blob_store);
        (spec, None)
    } else {
        // Ephemeral mode (original behavior).
        // Use MemoryStore to avoid JSONL writes — the gateway uses EphemeralSessionService
        // so agent-level persistence is not needed. This avoids failures on read-only
        // filesystems (e.g., GKE containers) where the default JSONL store can't write.
        let temp_dir = tempfile::tempdir().expect("create temp dir for agent working space");
        let binary_blob_store: Arc<dyn BinaryBlobStore> = Arc::new(ObjectStoreBlobStore::memory());
        let blob_store: Arc<dyn meerkat_core::BlobStore> =
            Arc::new(Base64BlobStoreAdapter::new(binary_blob_store.clone()));
        let runtime_store: Arc<dyn meerkat_runtime::RuntimeStore> =
            Arc::new(meerkat_runtime::InMemoryRuntimeStore::new());
        let adapter = Arc::new(meerkat_runtime::MeerkatMachine::persistent(
            Arc::clone(&runtime_store),
            Arc::clone(&blob_store),
        ));
        let mut factory = AgentFactory::new(temp_dir.path())
            .builtins(false)
            .comms(true)
            .session_store(Arc::new(meerkat::MemoryStore::new()));
        if image_generation {
            factory = factory.with_image_generation_machine(adapter.clone());
        }
        let mut inner_builder = FactoryAgentBuilder::new(factory, Config::default());
        inner_builder.default_blob_store = Some(blob_store);
        let callback_builder = StdioCallbackAgentBuilder {
            inner: inner_builder,
            bridge: bridge.clone(),
            has_session_builder,
            session_store: None,
        };
        let session_service = Arc::new(EphemeralSessionService::new(callback_builder, 16));

        let mut spec = MobBootstrapSpec::new(definition, MobStorage::in_memory(), session_service)
            .with_options(MobBootstrapOptions {
                allow_ephemeral_sessions: true,
                notify_orchestrator_on_resume: true,
                default_llm_client: None,
            });
        spec.runtime_adapter = Some(adapter);
        spec.binary_blob_store = Some(binary_blob_store);
        (spec, Some(temp_dir))
    };

    // Wire callback/after_create — notify Python/TS SDK after each session creation.
    // Uses notify_reliable to avoid silent drops under backpressure.
    let mob_spec = if has_session_builder {
        let after_bridge = bridge.clone();
        mob_spec.with_after_create_hook(Arc::new(
            move |session_id: meerkat_core::types::SessionId, ctx| {
                let b = after_bridge.clone();
                Box::pin(async move {
                    b.notify_reliable(
                        "callback/after_create",
                        json!({
                            "session_id": session_id.to_string(),
                            "model": ctx.model,
                            "labels": ctx.labels,
                            "system_prompt": ctx.system_prompt,
                        }),
                    )
                    .await;
                })
            },
        ))
    } else {
        mob_spec
    };

    let timeout = Duration::from_secs(30);
    let mut runtime = UnifiedRuntime::bootstrap(mob_spec, module_config, timeout)
        .await
        .unwrap_or_else(|e| {
            let error_response = json!({
                "jsonrpc": "2.0",
                "id": request_id,
                "error": { "code": -32603, "message": format!("Runtime bootstrap failed: {e}") }
            });
            let mut stdout = std::io::stdout().lock();
            let _ = writeln!(
                stdout,
                "{}",
                serde_json::to_string(&error_response)
                    .unwrap_or_else(|_| r#"{"error":"serialization failed"}"#.to_string())
            );
            let _ = stdout.flush();
            std::process::exit(1);
        });

    // 5b. Wire error hook — forwards ErrorEvents to Python as JSON-RPC notifications
    {
        let error_bridge = bridge.clone();
        runtime.set_error_hook(Arc::new(move |event| {
            let b = error_bridge.clone();
            Box::pin(async move {
                if let Ok(params) = serde_json::to_value(&event) {
                    b.notify("mobkit/on_error", params);
                }
            })
        }));
    }

    // 5c. Build identity-first runtime if providers are configured
    let identity_ctx: Option<meerkat_mobkit::rpc::IdentityFirstContext> = if has_roster_provider {
        use meerkat_mobkit::identity_first::{
            DurabilityPolicy, IdentityRuntime, IdentityRuntimeConfig, LocalContinuityStore,
            LocalLeaseProvider, RosterContext,
            contracts::LeaseProvider as LeaseProviderTrait,
            gateway_bridges::{
                GatewayAgentCustomizer, GatewayRosterProvider, GatewayTopologyProvider,
            },
        };

        // Build provider bridges
        let continuity_store: Arc<dyn meerkat_mobkit::identity_first::ContinuityStore> =
            if let Some(ref state_path) = persistent_state {
                Arc::new(
                    LocalContinuityStore::open(state_path.join("continuity.db")).unwrap_or_else(
                        |e| {
                            fail_init(
                                &request_id,
                                -32603,
                                format!("failed to open continuity store: {e}"),
                            );
                        },
                    ),
                )
            } else {
                Arc::new(
                    LocalContinuityStore::open(
                        std::env::temp_dir()
                            .join(format!("mobkit-continuity-{}.db", std::process::id())),
                    )
                    .unwrap_or_else(|e| {
                        fail_init(
                            &request_id,
                            -32603,
                            format!("failed to open continuity store: {e}"),
                        );
                    }),
                )
            };

        let lease_provider: Arc<dyn LeaseProviderTrait> = Arc::new(LocalLeaseProvider::new());

        // Construct the session bridge from the mob handle. The gateway
        // uses the raw bootstrap path (not UnifiedRuntimeBuilder), so
        // session_bridge() won't be set — build it directly.
        let mob_handle = runtime.mob_handle();
        let bridge_arc: Arc<dyn meerkat_mobkit::identity_first::SessionBridge> =
            if let Some(session_service) = runtime.mob_runtime().session_service().cloned() {
                Arc::new(
                    meerkat_mobkit::identity_first::MobSessionBridge::with_session_service(
                        mob_handle,
                        session_service,
                    ),
                )
            } else {
                Arc::new(meerkat_mobkit::identity_first::MobSessionBridge::new(
                    mob_handle,
                ))
            };

        let irt = IdentityRuntime::new(IdentityRuntimeConfig {
            continuity_store,
            lease_provider,
            runtime_instance_id: format!("gateway-{}", std::process::id()),
            has_runtime_store: persistent_state.is_some(),
            durability_policy: DurabilityPolicy::SyncWriteThrough,
            bridge: Some(bridge_arc),
            default_timeout: None,
        });

        // Build provider bridges for callbacks to Python
        let roster: Arc<dyn meerkat_mobkit::identity_first::contracts::RosterProvider> =
            Arc::new(GatewayRosterProvider::new(bridge.clone()));
        let topology: Option<Arc<dyn meerkat_mobkit::identity_first::contracts::TopologyProvider>> =
            if has_topology_provider {
                Some(Arc::new(GatewayTopologyProvider::new(bridge.clone())))
            } else {
                None
            };
        let customizer: Option<
            Arc<dyn meerkat_mobkit::identity_first::contracts::AgentCustomizer>,
        > = if has_agent_customizer {
            Some(Arc::new(GatewayAgentCustomizer::new(bridge.clone())))
        } else {
            None
        };

        // Call restore_flow to bootstrap identities from the roster provider
        let roster_specs = roster
            .roster(&RosterContext {
                mob_definition: None,
                previous_identities: Vec::new(),
            })
            .await
            .unwrap_or_else(|e| {
                fail_init(&request_id, -32603, format!("roster provider failed: {e}"));
            });

        if let Err(e) = meerkat_mobkit::identity_first::restore_flow(
            &irt,
            &roster_specs,
            topology.as_deref(),
            customizer.as_deref(),
        )
        .await
        {
            fail_init(
                &request_id,
                -32603,
                format!("identity-first restore_flow failed: {e}"),
            );
        }

        Some(meerkat_mobkit::rpc::IdentityFirstContext {
            runtime: Arc::new(irt),
            roster_provider: roster,
            topology_provider: topology,
            customizer,
        })
    } else {
        None
    };

    // 6. Bind HTTP server on ephemeral port
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
        .await
        .expect("bind ephemeral port");
    let port = listener.local_addr().expect("local addr").port();
    let http_base_url = format!("http://127.0.0.1:{port}");

    // 7. Start HTTP with graceful shutdown
    let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false);
    let app = runtime.build_reference_app_router(minimal_decision_state());
    let serve_task = tokio::spawn({
        let mut shutdown_rx = shutdown_rx.clone();
        async move {
            axum::serve(listener, app)
                .with_graceful_shutdown(async move {
                    shutdown_rx.changed().await.ok();
                })
                .await
        }
    });

    // 8. Send init response via stdout channel
    let loaded_modules = runtime.loaded_modules().await;
    let runtime_origin = if is_workspace_config {
        "workspace_config"
    } else {
        "fallback_minimal"
    };
    let runtime_fingerprint = {
        let mut hasher = Sha256::new();
        hasher.update(mob_config_toml.as_bytes());
        hasher.update(
            serde_json::to_string(&loaded_modules)
                .unwrap_or_default()
                .as_bytes(),
        );
        format!("{:x}", hasher.finalize())
    };
    let init_response = json!({
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "http_base_url": http_base_url,
            "loaded_modules": loaded_modules,
            "contract_version": MOBKIT_CONTRACT_VERSION,
            "runtime_origin": runtime_origin,
            "runtime_fingerprint": runtime_fingerprint,
        }
    });
    let _ = stdout_tx
        .send(
            serde_json::to_string(&init_response)
                .unwrap_or_else(|_| r#"{"error":"serialization failed"}"#.to_string()),
        )
        .await;

    // 9. RPC dispatch loop: process queued requests from the stdin reader task
    {
        loop {
            let request_line = tokio::select! {
                line = rpc_rx.recv() => match line {
                    Some(l) => l,
                    None => break, // stdin reader closed (EOF or error)
                },
                _ = tokio::signal::ctrl_c() => break,
            };
            let response = handle_unified_rpc_json(
                &runtime,
                &request_line,
                timeout,
                Some(&http_base_url),
                identity_ctx.as_ref(),
            )
            .await;
            if !response.is_empty() {
                let _ = stdout_tx.send(response).await;
            }
        }
    }

    // 10. Graceful shutdown: stop HTTP server, then runtime
    let _ = shutdown_tx.send(true);
    let _ = serve_task.await;
    runtime.shutdown().await;
    drop(rpc_tx);
    drop(stdout_tx);
    let _ = stdin_reader.await;
    let _ = stdout_writer.await;
    drop(_temp_dir);
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.iter().any(|a| a == "--persistent") {
        run_persistent();
    } else {
        run_single_shot();
    }
}
