//! Turn state, one stdin owner, typed events, and durable checkpoints.
use crate::{ChatMessage, Session};
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::{
    collections::VecDeque,
    future::Future,
    io::BufRead,
    path::PathBuf,
    sync::{Arc, Mutex},
    time::Duration,
};
use tokio::sync::{oneshot, watch, Notify};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum StepStatus {
    Pending,
    InProgress,
    Completed,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct PlanStep {
    pub step: String,
    pub status: StepStatus,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentEvent {
    Status {
        message: String,
    },
    Progress {
        message: String,
    },
    PlanUpdated {
        plan: Vec<PlanStep>,
        explanation: Option<String>,
    },
    ToolStarted {
        id: String,
        name: String,
    },
    ToolFinished {
        id: String,
        name: String,
        success: bool,
    },
    Waiting {
        stage: String,
        elapsed_secs: u64,
    },
    Question {
        question: Value,
    },
    Steering {
        message: String,
    },
    TurnEnded {
        verdict: String,
        message: String,
    },
}

#[derive(Default)]
struct InputState {
    lines: VecDeque<String>,
    steering: VecDeque<String>,
    question: Option<oneshot::Sender<Option<String>>>,
    eof: bool,
    running: bool,
}
pub struct InputBroker {
    state: Mutex<InputState>,
    notify: Notify,
    cancel: watch::Sender<bool>,
    persistent: bool,
}
impl std::fmt::Debug for InputBroker {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("InputBroker")
    }
}
impl InputBroker {
    pub fn new() -> Arc<Self> {
        Self::with_persistence(false)
    }
    pub(crate) fn persistent() -> Arc<Self> {
        Self::with_persistence(true)
    }
    fn with_persistence(persistent: bool) -> Arc<Self> {
        Arc::new(Self {
            persistent,
            state: Mutex::new(InputState::default()),
            notify: Notify::new(),
            cancel: watch::channel(false).0,
        })
    }
    pub fn stdin() -> Arc<Self> {
        let broker = Self::new();
        let reader = broker.clone();
        // This is the only stdin reader for run/agent, including human tools.
        std::thread::spawn(move || {
            for line in std::io::stdin().lock().lines() {
                match line {
                    Ok(line) => reader.feed(line),
                    Err(_) => break,
                }
            }
            reader.close();
        });
        broker
    }
    pub(crate) fn can_ask(&self) -> bool {
        !self.state.lock().unwrap().eof
    }
    pub(crate) fn has_question(&self) -> bool {
        self.state.lock().unwrap().question.is_some()
    }
    pub fn feed(&self, line: String) {
        let mut state = self.state.lock().unwrap();
        if state.running && matches!(line.trim(), "/stop" | "/exit" | "/quit") {
            self.cancel.send_replace(true);
            return;
        }
        if let Some(answer) = state.question.take() {
            let _ = answer.send(Some(line));
        } else if state.running {
            state.steering.push_back(line);
        } else {
            state.lines.push_back(line);
        }
        drop(state);
        self.notify.notify_one();
    }
    pub fn close(&self) {
        let mut state = self.state.lock().unwrap();
        state.eof = true;
        if let Some(answer) = state.question.take() {
            let _ = answer.send(None);
        }
        drop(state);
        self.notify.notify_one();
    }
    pub fn begin(&self) {
        self.cancel.send_replace(false);
        self.state.lock().unwrap().running = true;
    }
    pub fn end(&self) {
        let mut s = self.state.lock().unwrap();
        s.running = false;
        s.question.take();
        let rest = s.steering.drain(..).collect::<Vec<_>>();
        s.lines.extend(rest);
    }
    pub fn interrupt(&self) {
        if self.state.lock().unwrap().running {
            self.cancel.send_replace(true);
        } else {
            self.feed("/exit".into());
        }
    }
    pub fn cancelled(&self) -> bool {
        *self.cancel.borrow()
    }
    pub async fn cancellation(&self) {
        let mut rx = self.cancel.subscribe();
        loop {
            if *rx.borrow_and_update() {
                return;
            }
            if rx.changed().await.is_err() {
                return;
            }
        }
    }
    pub fn drain(&self) -> Vec<String> {
        self.state
            .lock()
            .unwrap()
            .steering
            .drain(..)
            .filter(|s| !s.trim().is_empty())
            .collect()
    }
    pub async fn next(&self) -> Option<String> {
        loop {
            let notified = self.notify.notified();
            {
                let mut s = self.state.lock().unwrap();
                if let Some(line) = s.lines.pop_front() {
                    return Some(line);
                }
                if s.eof {
                    return None;
                }
            }
            notified.await;
        }
    }
    pub async fn ask(&self, seconds: u64, on_ready: impl FnOnce()) -> Result<String> {
        let (tx, rx) = oneshot::channel();
        {
            let mut s = self.state.lock().unwrap();
            if s.eof {
                drop(s);
                on_ready();
                return Err(anyhow!("stdin closed; answer in a resumed session"));
            }
            s.question = Some(tx);
        }
        on_ready();
        let result = tokio::select! {
            biased;
            _ = self.cancellation() => Err(anyhow!("cancelled")),
            answer = async { if self.persistent { Ok(rx.await) } else { tokio::time::timeout(Duration::from_secs(seconds.clamp(1, 3600)), rx).await } } => match answer {
                Ok(Ok(Some(line))) => Ok(line),
                Ok(_) => Err(anyhow!("stdin closed")),
                Err(_) => Err(anyhow!("user input timed out")),
            }
        };
        self.state.lock().unwrap().question.take();
        result
    }
}

type CheckpointSink = Arc<dyn Fn(&Session) -> Result<()> + Send + Sync>;
type EventSink = Arc<dyn Fn(&AgentEvent) + Send + Sync>;

#[derive(Default)]
pub struct Runtime {
    pub input: Option<Arc<InputBroker>>,
    pub plan: Mutex<Vec<PlanStep>>,
    pub pending_question: Mutex<Option<Value>>,
    pub active_command: Mutex<Option<Value>>,
    pub browser_url: Mutex<Option<String>>,
    pub(crate) browser_profile: Mutex<Option<PathBuf>>,
    checkpoint: Mutex<Option<Session>>,
    pub events: Mutex<Vec<AgentEvent>>,
    event_sink: Mutex<Option<EventSink>>,
    checkpoint_sink: Mutex<Option<CheckpointSink>>,
}
impl std::fmt::Debug for Runtime {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("Runtime")
    }
}
impl Runtime {
    pub fn with_input(input: Arc<InputBroker>) -> Arc<Self> {
        Arc::new(Self {
            input: Some(input),
            ..Default::default()
        })
    }
    pub fn for_session(session: &Session, input: Arc<InputBroker>) -> Arc<Self> {
        Arc::new(Self {
            input: Some(input),
            plan: Mutex::new(session.plan.clone()),
            pending_question: Mutex::new(session.pending_question.clone()),
            active_command: Mutex::new(session.active_command.clone()),
            browser_url: Mutex::new(session.last_browser_url.clone()),
            browser_profile: Mutex::new(None),
            checkpoint: Mutex::new(Some(session.clone())),
            events: Mutex::new(vec![]),
            event_sink: Mutex::new(None),
            checkpoint_sink: Mutex::new(None),
        })
    }
    pub(crate) fn set_checkpoint_handler(
        &self,
        handler: impl Fn(&Session) -> Result<()> + Send + Sync + 'static,
    ) {
        *self.checkpoint_sink.lock().unwrap() = Some(Arc::new(handler));
    }
    fn save_checkpoint(&self, session: &Session) -> Result<()> {
        if let Some(sink) = self.checkpoint_sink.lock().unwrap().clone() {
            sink(session)
        } else {
            crate::save_session(session).map(|_| ())
        }
    }
    pub fn set_event_handler(&self, handler: impl Fn(&AgentEvent) + Send + Sync + 'static) {
        *self.event_sink.lock().unwrap() = Some(Arc::new(handler));
    }
    pub fn emit(&self, event: AgentEvent) {
        let sink = self.event_sink.lock().unwrap().clone();
        if let Some(sink) = sink {
            sink(&event);
        }
        self.events.lock().unwrap().push(event);
    }
    pub async fn wait<T>(&self, stage: &str, future: impl Future<Output = Result<T>>) -> Result<T> {
        tokio::pin!(future);
        let start = std::time::Instant::now();
        let mut tick = tokio::time::interval(Duration::from_secs(20));
        tick.tick().await;
        loop {
            tokio::select! {
                biased;
                _ = async { match &self.input { Some(i) => i.cancellation().await, None => std::future::pending().await } } => return Err(anyhow!("cancelled")),
                result = &mut future => return result,
                _ = tick.tick() => self.emit(AgentEvent::Waiting { stage: stage.into(), elapsed_secs: start.elapsed().as_secs() }),
            }
        }
    }
    pub fn cancelled(&self) -> bool {
        self.input.as_ref().is_some_and(|i| i.cancelled())
    }
    pub fn steering(&self) -> Vec<String> {
        self.input.as_ref().map(|i| i.drain()).unwrap_or_default()
    }
    pub fn update_plan(&self, args: &Value) -> Result<Value> {
        let next: Vec<PlanStep> = serde_json::from_value(args["plan"].clone())?;
        if next.is_empty()
            || next.len() > 12
            || next.iter().any(|s| s.step.trim().is_empty())
            || next
                .iter()
                .filter(|s| s.status == StepStatus::InProgress)
                .count()
                > 1
        {
            return Err(anyhow!(
                "plan requires 1–12 nonempty steps, at most one in_progress"
            ));
        }
        let mut current = self.plan.lock().unwrap();
        let explanation = args["explanation"]
            .as_str()
            .filter(|s| !s.trim().is_empty())
            .map(str::to_string);
        if !current.is_empty()
            && current.iter().map(|s| &s.step).collect::<Vec<_>>()
                != next.iter().map(|s| &s.step).collect::<Vec<_>>()
            && explanation.is_none()
        {
            return Err(anyhow!("explain why the plan steps changed"));
        }
        *current = next.clone();
        drop(current);
        self.emit(AgentEvent::PlanUpdated {
            plan: next,
            explanation,
        });
        Ok(json!({"updated":true}))
    }
    pub fn unfinished(&self) -> bool {
        self.plan
            .lock()
            .unwrap()
            .iter()
            .any(|s| s.status != StepStatus::Completed)
    }
    pub fn checkpoint(&self, messages: &[ChatMessage], pending_tool: Option<&str>) -> Result<()> {
        let mut guard = self.checkpoint.lock().unwrap();
        if let Some(s) = guard.as_mut() {
            s.messages = messages.to_vec();
            s.plan = self.plan.lock().unwrap().clone();
            s.pending_question = self.pending_question.lock().unwrap().clone();
            s.pending_tool = pending_tool.map(str::to_string);
            s.active_command = self.active_command.lock().unwrap().clone();
            s.last_browser_url = self
                .browser_url
                .lock()
                .unwrap()
                .clone()
                .or(s.last_browser_url.clone());
            s.touch();
            self.save_checkpoint(s)?;
        }
        Ok(())
    }
    pub fn sync_session(&self, session: &mut Session) {
        session.plan = self.plan.lock().unwrap().clone();
        session.pending_question = self.pending_question.lock().unwrap().clone();
        session.pending_tool = None;
        session.active_command = self.active_command.lock().unwrap().clone();
    }
    pub fn profile_dir(&self) -> Option<PathBuf> {
        if let Some(path) = self.browser_profile.lock().unwrap().clone() {
            return Some(path);
        }
        self.checkpoint.lock().unwrap().as_ref().and_then(|s| {
            (if self.checkpoint_sink.lock().unwrap().is_some() {
                Ok(crate::team::data_dir().join("profiles"))
            } else {
                crate::sessions_dir()
            })
            .ok()
            .map(|d| d.join(format!("{}.browser", s.id)))
        })
    }
    pub fn save_output(&self, output: &str) -> Result<Option<String>> {
        if output.len() <= 16000 {
            return Ok(None);
        }
        let guard = self.checkpoint.lock().unwrap();
        let Some(s) = guard.as_ref() else {
            return Ok(None);
        };
        // Artifacts belong to the workspace so read_file can access them.
        let dir = s.cwd.join(".grokboy-output").join(&s.id);
        std::fs::create_dir_all(&dir)?;
        let path = dir.join(format!("{}.txt", uuid::Uuid::new_v4()));
        std::fs::write(&path, output)?;
        Ok(Some(json!({"preview":output.chars().take(4000).collect::<String>(),"output_file":path,"bytes":output.len(),"truncated":true,"hint":"read_file with offset/limit to inspect full output"}).to_string()))
    }
    pub async fn question(&self, args: &Value) -> Result<Value> {
        let input = self
            .input
            .as_ref()
            .ok_or_else(|| anyhow!("no interactive input channel"))?;
        *self.pending_question.lock().unwrap() = Some(args.clone());
        // Persist pending question against latest pre-tool checkpoint.
        {
            let mut guard = self.checkpoint.lock().unwrap();
            if let Some(s) = guard.as_mut() {
                s.pending_question = Some(args.clone());
                s.last_browser_url = self
                    .browser_url
                    .lock()
                    .unwrap()
                    .clone()
                    .or(s.last_browser_url.clone());
                self.save_checkpoint(s)?;
            }
        }
        let line = input
            .ask(args["timeout_secs"].as_u64().unwrap_or(300), || {
                self.emit(AgentEvent::Question {
                    question: args.clone(),
                })
            })
            .await?;
        *self.pending_question.lock().unwrap() = None;
        let answer = line
            .trim()
            .parse::<usize>()
            .ok()
            .and_then(|n| n.checked_sub(1))
            .or_else(|| {
                let text = line.trim().to_ascii_lowercase();
                if text.len() == 1 {
                    text.as_bytes()[0]
                        .checked_sub(b'a')
                        .filter(|n| *n < 26)
                        .map(usize::from)
                } else {
                    None
                }
            })
            .and_then(|n| args["options"].get(n))
            .and_then(Value::as_str)
            .unwrap_or(line.trim())
            .to_string();
        Ok(json!({"answer":answer}))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[tokio::test]
    async fn queued_tasks_steering_and_answers_do_not_steal_each_other() {
        let input = InputBroker::new();
        input.feed("task one".into());
        input.feed("task two".into());
        assert_eq!(input.next().await.as_deref(), Some("task one"));
        input.begin();
        input.feed("keep original goal, add constraint".into());
        assert_eq!(input.drain(), vec!["keep original goal, add constraint"]);
        let answer = input.ask(2, || input.feed("answer".into())).await.unwrap();
        assert_eq!(answer, "answer");
        assert!(input.drain().is_empty());
        input.end();
        assert_eq!(input.next().await.as_deref(), Some("task two"));
    }
    #[tokio::test]
    async fn cancellation_interrupts_a_pending_question_without_losing_next_input() {
        let input = InputBroker::new();
        input.begin();
        assert!(input.ask(30, || input.interrupt()).await.is_err());
        input.end();
        input.feed("resume".into());
        assert_eq!(input.next().await.as_deref(), Some("resume"));
    }
    #[test]
    fn plan_changes_require_explanation_and_one_active_step() {
        let runtime = Runtime::default();
        runtime
            .update_plan(&json!({"plan":[{"step":"inspect","status":"in_progress"}]}))
            .unwrap();
        assert!(runtime
            .update_plan(&json!({"plan":[{"step":"other","status":"in_progress"}]}))
            .is_err());
        assert!(runtime.update_plan(&json!({"explanation":"new discovery","plan":[{"step":"other","status":"in_progress"},{"step":"third","status":"in_progress"}]})).is_err());
        runtime.update_plan(&json!({"explanation":"new discovery","plan":[{"step":"other","status":"completed"}]})).unwrap();
        assert!(!runtime.unfinished());
    }
}