//! Long-running claude process for the shard. Spawned once, fed turns via
//! stdin (stream-json), emits events via stdout. Survives across many turns
//! within configured limits, then respawns fresh.

use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::{Duration, Instant};

use anyhow::{Context, bail};
use matrix_sdk::ruma::OwnedRoomId;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, Command};
use tokio::sync::mpsc;

use crate::paths;

/// Curated system-prompt preamble. Replaces Claude Code's default (which
/// assumes stdout reaches the user, software-engineering framing, etc.).
/// Followed by SYSTEM.md content at runtime. Source: `prompts/shard_preamble.md`.
const SYSTEM_PROMPT_PREAMBLE: &str = include_str!("../prompts/shard_preamble.md");

/// Permissions JSON passed via --settings. Without this, WebFetch and
/// WebSearch fall back to interactive permission prompts, which the daemon
/// has no way to answer (claude is a child process). Allowing them broadly
/// is fine: the lab is sandboxed, and the tools are read-only HTTP GET.
const SHARD_SETTINGS_JSON: &str = r#"{"permissions":{"allow":["WebFetch","WebSearch"]}}"#;

/// One stream-json event from claude's stdout. Only fields we care about.
#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
pub enum StreamEvent {
    #[serde(rename = "system")]
    System {
        #[serde(default)]
        subtype: String,
        #[serde(default)]
        session_id: Option<String>,
    },
    #[serde(rename = "assistant")]
    Assistant {
        #[serde(default)]
        message: serde_json::Value,
    },
    #[serde(rename = "result")]
    Result {
        #[serde(default)]
        subtype: String,
        #[serde(default)]
        is_error: bool,
        #[serde(default)]
        result: String,
        #[serde(default)]
        stop_reason: Option<String>,
    },
    #[serde(rename = "rate_limit_event")]
    RateLimit {},
    #[serde(other)]
    Other,
}

#[derive(Debug, Serialize)]
struct UserMessage<'a> {
    #[serde(rename = "type")]
    kind: &'static str,
    message: UserMessageBody<'a>,
}

#[derive(Debug, Serialize)]
struct UserMessageBody<'a> {
    role: &'static str,
    content: Vec<UserContent<'a>>,
}

#[derive(Debug, Serialize)]
struct UserContent<'a> {
    #[serde(rename = "type")]
    kind: &'static str,
    text: &'a str,
}

/// A live claude shard process. Owns stdin and an mpsc receiver for stdout
/// events. When it exits or is dropped, the child is cleaned up.
pub struct ShardSession {
    child: Child,
    stdin: ChildStdin,
    events: mpsc::Receiver<StreamEvent>,
    /// When the session was spawned.
    started: Instant,
    /// Last time a turn finished.
    last_used: Instant,
    /// Number of turns processed.
    pub turn_count: u32,
    /// Rooms we've sent at least one turn for in this session. Used to decide
    /// whether to include `previously_seen` context in a turn.
    pub rooms_seen: HashSet<OwnedRoomId>,
    /// Mtimes of identity/CHANGELOG files at session start - for refresh.
    mtime_snapshot: Vec<(PathBuf, std::time::SystemTime)>,
}

pub struct SpawnConfig<'a> {
    pub model: &'a str,
    pub mcp_config_path: &'a Path,
    pub allowed_tools: &'a str,
}

impl ShardSession {
    pub async fn spawn(cfg: SpawnConfig<'_>) -> anyhow::Result<Self> {
        let identity_dir = paths::identity_dir();

        // Build the shard's system prompt from a curated preamble (replaces
        // Claude Code's default, which is wrong for chat use - it assumes
        // stdout reaches the user, software-engineering framing, etc.) plus
        // SYSTEM.md (the harness contract).
        let system_md_path = identity_dir.join("SYSTEM.md");
        let system_md = tokio::fs::read_to_string(&system_md_path)
            .await
            .with_context(|| format!("reading {}", system_md_path.display()))?;
        let system_prompt = format!("{SYSTEM_PROMPT_PREAMBLE}\n{system_md}");

        let mut cmd = Command::new("claude");
        cmd.args([
            "--print",
            "--verbose",
            "--input-format",
            "stream-json",
            "--output-format",
            "stream-json",
            "--model",
            cfg.model,
            "--add-dir",
            &identity_dir.to_string_lossy(),
            "--allowedTools",
            cfg.allowed_tools,
            "--mcp-config",
            &cfg.mcp_config_path.to_string_lossy(),
            "--system-prompt",
            &system_prompt,
            "--settings",
            SHARD_SETTINGS_JSON,
        ]);
        cmd.current_dir(&identity_dir);
        cmd.stdin(Stdio::piped());
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
        cmd.kill_on_drop(true);

        let mut child = cmd.spawn().context("spawn claude")?;
        let stdin = child.stdin.take().context("claude stdin missing")?;
        let stdout = child.stdout.take().context("claude stdout missing")?;
        let stderr = child.stderr.take().context("claude stderr missing")?;

        let (tx, rx) = mpsc::channel(256);

        // stdout reader: parse stream-json line by line, push events
        let tx_out = tx.clone();
        tokio::spawn(async move {
            let mut lines = BufReader::new(stdout).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                if line.trim().is_empty() {
                    continue;
                }
                let ev: StreamEvent = match serde_json::from_str(&line) {
                    Ok(v) => v,
                    Err(e) => {
                        tracing::debug!(?e, line = %line, "shard: failed to parse stream-json line");
                        continue;
                    }
                };
                if tx_out.send(ev).await.is_err() {
                    break;
                }
            }
            tracing::info!("shard: stdout stream closed");
        });

        // stderr drainer: log claude's stderr at warn level
        tokio::spawn(async move {
            let mut lines = BufReader::new(stderr).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                tracing::warn!("claude stderr: {line}");
            }
        });

        let mtime_snapshot = snapshot_identity_mtimes();
        tracing::info!("shard: spawned");

        Ok(Self {
            child,
            stdin,
            events: rx,
            started: Instant::now(),
            last_used: Instant::now(),
            turn_count: 0,
            rooms_seen: HashSet::new(),
            mtime_snapshot,
        })
    }

    /// Send one user-turn JSON to claude and wait for the next `result`.
    /// Returns the result text (claude's final assistant output).
    pub async fn run_turn(&mut self, turn_body: &str) -> anyhow::Result<String> {
        let msg = UserMessage {
            kind: "user",
            message: UserMessageBody {
                role: "user",
                content: vec![UserContent {
                    kind: "text",
                    text: turn_body,
                }],
            },
        };
        let mut json = serde_json::to_string(&msg)?;
        json.push('\n');
        self.stdin.write_all(json.as_bytes()).await?;
        self.stdin.flush().await?;

        // Drain events until we hit a `result` (turn end).
        loop {
            let ev = self
                .events
                .recv()
                .await
                .context("shard: stdout closed before result")?;
            match ev {
                StreamEvent::System { .. } | StreamEvent::RateLimit {} | StreamEvent::Other => {}
                StreamEvent::Assistant { message } => {
                    log_assistant_text(&message);
                }
                StreamEvent::Result {
                    is_error,
                    result,
                    stop_reason,
                    ..
                } => {
                    self.last_used = Instant::now();
                    self.turn_count += 1;
                    if is_error {
                        bail!("turn ended with is_error=true: {result}");
                    }
                    tracing::info!(
                        turn = self.turn_count,
                        stop_reason = ?stop_reason,
                        "shard: turn complete"
                    );
                    return Ok(result);
                }
            }
        }
    }

    pub fn should_refresh(&mut self, idle: Duration, max_turns: u32) -> Option<&'static str> {
        if self.last_used.elapsed() > idle {
            return Some("idle gap exceeded");
        }
        if self.turn_count >= max_turns {
            return Some("max turn count reached");
        }
        if mtimes_changed(&self.mtime_snapshot) {
            return Some("identity files changed");
        }
        if self.child.try_wait().ok().flatten().is_some() {
            return Some("child process exited");
        }
        None
    }

    pub async fn shutdown(mut self) {
        // Close stdin → claude exits gracefully on EOF
        drop(self.stdin);
        // Bounded wait for clean exit
        let _ = tokio::time::timeout(Duration::from_secs(3), self.child.wait()).await;
        let _ = self.child.kill().await;
        tracing::info!(
            turns = self.turn_count,
            uptime = ?self.started.elapsed(),
            "shard: session ended"
        );
    }
}

fn log_assistant_text(message: &serde_json::Value) {
    let Some(content) = message.get("content").and_then(|c| c.as_array()) else {
        return;
    };
    for item in content {
        let Some(kind) = item.get("type").and_then(|t| t.as_str()) else {
            continue;
        };
        match kind {
            "text" => {
                if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
                    let preview: String = text.chars().take(200).collect();
                    tracing::info!("shard text: {preview}");
                    tracing::trace!("shard full text: {text}");
                }
            }
            "thinking" => {
                if let Some(text) = item.get("thinking").and_then(|t| t.as_str()) {
                    tracing::trace!("shard thinking: {text}");
                }
            }
            _ => {}
        }
    }
}

fn snapshot_identity_mtimes() -> Vec<(PathBuf, std::time::SystemTime)> {
    let id = paths::identity_dir();
    let state = paths::state_dir();
    let candidates = [
        id.join("CLAUDE.md"),
        id.join("SYSTEM.md"),
        id.join("notes.md"),
        state.join("CHANGELOG.md"),
    ];
    candidates
        .into_iter()
        .filter_map(|p| std::fs::metadata(&p).and_then(|m| m.modified()).ok().map(|t| (p, t)))
        .collect()
}

fn mtimes_changed(snapshot: &[(PathBuf, std::time::SystemTime)]) -> bool {
    for (path, prev) in snapshot {
        if let Ok(now) = std::fs::metadata(path).and_then(|m| m.modified()) {
            if &now != prev {
                return true;
            }
        }
    }
    false
}