//! Deciding finished turns when `loopd` reports one (M5c spec, section 5). At each report every
//! session with undecided turns is walked: its log is read, `brokerd` is asked once for its state,
//! and each undecided turn is decided and appended to the ledger at once. A turn that cannot be
//! decided now (`brokerd` silent, the log unreadable) waits for the next report of any session.

use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, MutexGuard};

use bxcore::{LogRecord, SessionId, SessionStateReply, Timestamp};

use crate::honcho::decide::{decide, turn_records, turns};
use crate::honcho::ledger::{self, Item, Ledger, Record};
use crate::honcho::state;
use crate::honcho::status::{self, Shared};

/// The ledger's file name, in `memory/`.
pub const LEDGER: &str = "honcho.jsonl";

/// The runbook entry for a session log with a damaged line.
pub const LOG_RUNBOOK: &str = "see docs/runbook.md#session-log-damaged";

/// `brokerd`'s answer for a session, or why there is none.
pub type Ask<'a> = &'a dyn Fn(&SessionId) -> Result<SessionStateReply, String>;

/// How one session's walk ended.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Walk {
    /// Every turn up to the reported one is decided: the session leaves the map.
    Done,
    /// Some turn waits: the session stays in the map.
    Waiting,
    /// `brokerd` did not answer, or the ledger could not be written: the walk stops.
    Stop,
}

/// The turns reported and not yet decided, and the ledger their decisions go to. The ledger has a
/// lock of its own, held only to read or append: a walk holds `TurnSync`'s lock while it asks
/// `brokerd`, and a memory write, decided under the index lock, takes only the ledger's.
pub struct TurnSync {
    ledger: Arc<Mutex<Ledger>>,
    agent: String,
    sessions: PathBuf,
    /// Each session with undecided turns, by its id, and the highest turn `loopd` reported for
    /// it. `SessionId` has no order, so the map is keyed by its text.
    reported: BTreeMap<String, (SessionId, u64)>,
    /// Whether the last question to `brokerd` went unanswered: an outage is logged once.
    broker_down: bool,
    /// Sessions whose log last had a damaged line: logged once each.
    unreadable: BTreeSet<String>,
    /// What `bxctl memory sync` shows beyond the ledger, shared with the worker and the report.
    status: Shared,
}

/// The path of `session`'s epoch-0 log under `sessions`. Later epochs are sibling files;
/// `read_session` concatenates them. Existing tests pin this path.
pub fn log_path(sessions: &Path, session: &SessionId) -> PathBuf {
    sessions.join(session.as_str()).join("0.jsonl")
}

/// Why a session log did not read. Each names the file; `Display` adds the runbook pointer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LogError {
    /// The file exists but could not be read (a permission, an I/O error): it may read later.
    Io(String),
    /// A complete line that does not parse, and its number: it will not read later.
    Damaged(String),
}

impl std::fmt::Display for LogError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LogError::Io(why) | LogError::Damaged(why) => write!(f, "{why}; {LOG_RUNBOOK}"),
        }
    }
}

impl std::error::Error for LogError {}

/// A session log's complete lines as records. A missing file has none. A last line without its
/// newline is one `loopd` is still writing, and is left out. A complete line that does not parse
/// is an error naming the file and the line, with the runbook pointer.
pub fn read_log(path: &Path) -> Result<Vec<LogRecord>, String> {
    read_records(path).map_err(|e| e.to_string())
}

/// `read_log`, telling a file that could not be read from a damaged line.
pub fn read_records(path: &Path) -> Result<Vec<LogRecord>, LogError> {
    let text = match std::fs::read(path) {
        Ok(bytes) => bytes,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(e) => return Err(LogError::Io(format!("{}: {e}", path.display()))),
    };
    let whole = match text.iter().rposition(|b| *b == b'\n') {
        Some(last) => last + 1,
        None => 0,
    };
    let mut records = Vec::new();
    for (index, line) in text
        .get(..whole)
        .unwrap_or_default()
        .split_inclusive(|b| *b == b'\n')
        .enumerate()
    {
        match serde_json::from_slice::<LogRecord>(line) {
            Ok(record) => records.push(record),
            Err(e) => {
                return Err(LogError::Damaged(format!(
                    "{}:{}: {e}",
                    path.display(),
                    index + 1
                )));
            }
        }
    }
    Ok(records)
}

/// Every `<n>.jsonl` of `session`, sorted by `n`. A missing directory is empty.
pub fn epoch_logs(sessions: &Path, session: &SessionId) -> Result<Vec<PathBuf>, LogError> {
    let dir = sessions.join(session.as_str());
    let entries = match std::fs::read_dir(&dir) {
        Ok(e) => e,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(e) => return Err(LogError::Io(format!("{}: {e}", dir.display()))),
    };
    let mut found: Vec<(u32, PathBuf)> = Vec::new();
    for entry in entries {
        let entry = match entry {
            Ok(e) => e,
            Err(e) => return Err(LogError::Io(format!("{}: {e}", dir.display()))),
        };
        let name = entry.file_name();
        let Some(name) = name.to_str() else { continue };
        let Some(stem) = name.strip_suffix(".jsonl") else {
            continue;
        };
        let Ok(n) = stem.parse::<u32>() else { continue };
        found.push((n, entry.path()));
    }
    found.sort_by_key(|(n, _)| *n);
    Ok(found.into_iter().map(|(_, p)| p).collect())
}

/// Every epoch's records, in epoch order. A missing session has none.
pub fn read_session(sessions: &Path, session: &SessionId) -> Result<Vec<LogRecord>, LogError> {
    let mut out = Vec::new();
    for path in epoch_logs(sessions, session)? {
        out.extend(read_records(&path)?);
    }
    Ok(out)
}

/// Whether a turn has ended with an answer: its last `Assistant` record asks for no tool.
pub fn finished(turn: &[LogRecord]) -> bool {
    let last = turn
        .iter()
        .rev()
        .find(|r| matches!(r, LogRecord::Assistant { .. }));
    matches!(last, Some(LogRecord::Assistant { tool_calls, .. }) if tool_calls.is_empty())
}

impl TurnSync {
    /// `ledger` read through at start; `agent` is `[honcho] agent`; `sessions` the directory of
    /// `loopd`'s session logs.
    pub fn new(ledger: Ledger, agent: &str, sessions: &Path) -> TurnSync {
        TurnSync {
            ledger: Arc::new(Mutex::new(ledger)),
            agent: agent.to_string(),
            sessions: sessions.to_path_buf(),
            reported: BTreeMap::new(),
            broker_down: false,
            unreadable: BTreeSet::new(),
            status: Shared::default(),
        }
    }

    /// The ledger, locked. Never hold it while taking another lock.
    pub fn ledger(&self) -> MutexGuard<'_, Ledger> {
        self.ledger.lock().unwrap_or_else(|p| p.into_inner())
    }

    /// The ledger's lock itself, for the worker and for memory writes.
    pub fn shared_ledger(&self) -> Arc<Mutex<Ledger>> {
        Arc::clone(&self.ledger)
    }

    /// The status `bxctl memory sync` shows, for the worker and the report.
    pub fn shared_status(&self) -> Shared {
        Arc::clone(&self.status)
    }

    /// The directory of `loopd`'s session logs.
    pub fn sessions(&self) -> &Path {
        &self.sessions
    }

    /// The sessions that still have undecided turns, with the highest turn reported for each.
    pub fn waiting(&self) -> Vec<(SessionId, u64)> {
        self.reported.values().cloned().collect()
    }

    /// After a restart: each session with a decided turn in the ledger waits again, at the turns
    /// its log holds, so that a turn left undecided before the restart is decided at the next
    /// walk (`decidable` holds back one still running). A session with no decided turn is left
    /// to its next report: sessions from before Honcho was turned on are never walked.
    pub fn resume(&mut self) {
        let mut sessions: Vec<SessionId> = Vec::new();
        for (item, _) in self.ledger().decisions() {
            if let Item::Turn { session, .. } = item
                && !sessions.contains(session)
            {
                sessions.push(session.clone());
            }
        }
        for session in sessions {
            if let Ok(records) = read_session(&self.sessions, &session) {
                self.reported.insert(
                    session.as_str().to_string(),
                    (session.clone(), turns(&records)),
                );
            }
        }
    }

    /// Append `record` to the ledger (the worker's `Sent` and `Failed`).
    pub fn append(&mut self, record: Record) -> Result<(), ledger::LedgerError> {
        self.ledger().append(record)
    }

    /// `loopd` reported that turn `turn` of `session` ended. Remember it, then walk.
    pub fn updated(
        &mut self,
        session: &SessionId,
        turn: u64,
        ask: Ask<'_>,
        now: Timestamp,
        log: &dyn Fn(&str),
    ) {
        let entry = self
            .reported
            .entry(session.as_str().to_string())
            .or_insert((session.clone(), 0));
        entry.1 = entry.1.max(turn);
        self.retry(ask, now, log);
    }

    /// Walk every session with undecided turns, in the order of their ids, until one stops the
    /// walk: at a report, and from the worker when it has nothing to send (M5c spec, section 5).
    pub fn retry(&mut self, ask: Ask<'_>, now: Timestamp, log: &dyn Fn(&str)) {
        for (session, turn) in self.waiting() {
            match self.walk(&session, turn, ask, now, log) {
                Walk::Done => {
                    self.reported.remove(session.as_str());
                }
                Walk::Waiting => {}
                Walk::Stop => return,
            }
        }
    }

    /// Decide `session`'s undecided turns from 1 to `reported`, as far as its log goes. The log's
    /// last turn waits until it has an answer: it may still be running.
    fn walk(
        &mut self,
        session: &SessionId,
        reported: u64,
        ask: Ask<'_>,
        now: Timestamp,
        log: &dyn Fn(&str),
    ) -> Walk {
        let Some(records) = self.records(session, log) else {
            return Walk::Waiting;
        };
        let (limit, held) = decidable(&records, reported);
        let item = |turn: u64| Item::Turn {
            session: session.clone(),
            turn,
        };
        let undecided: Vec<u64> = (1..=limit)
            .filter(|turn| !self.ledger().is_decided(&item(*turn)))
            .collect();
        if undecided.is_empty() {
            return if held { Walk::Waiting } else { Walk::Done };
        }
        let Some(state) = self.state(session, ask, now, log) else {
            return Walk::Stop;
        };
        for turn in undecided {
            let began = self.ledger().began();
            let Ok(action) = decide(&state, &records, turn, &self.agent, began) else {
                continue;
            };
            let record = Record::Decided {
                time: now,
                item: item(turn),
                action,
            };
            let appended = self.ledger().append(record);
            if let Err(e) = appended {
                log(&format!("memoryd: {e}; {}", ledger::RUNBOOK));
                return Walk::Stop;
            }
        }
        if held { Walk::Waiting } else { Walk::Done }
    }

    /// `session`'s log, read. A log that does not read is logged once, until it reads again.
    fn records(&mut self, session: &SessionId, log: &dyn Fn(&str)) -> Option<Vec<LogRecord>> {
        match read_session(&self.sessions, session).map_err(|e| e.to_string()) {
            Ok(records) => {
                self.unreadable.remove(session.as_str());
                Some(records)
            }
            Err(e) => {
                if self.unreadable.insert(session.as_str().to_string()) {
                    log(&format!("memoryd: {e}"));
                }
                None
            }
        }
    }

    /// `brokerd`'s answer for `session`. The first failure after an answer is logged, and so is
    /// the first answer after a failure.
    fn state(
        &mut self,
        session: &SessionId,
        ask: Ask<'_>,
        now: Timestamp,
        log: &dyn Fn(&str),
    ) -> Option<SessionStateReply> {
        match ask(session) {
            Ok(state) => {
                if self.broker_down {
                    self.broker_down = false;
                    status::update(&self.status, |s| s.broker_down_since = None);
                    log("memoryd: brokerd answers on its state socket again; turns are decided");
                }
                Some(state)
            }
            Err(e) => {
                if !self.broker_down {
                    self.broker_down = true;
                    status::update(&self.status, |s| s.broker_down_since = Some(now));
                    log(&format!(
                        "memoryd: {e}; turns wait undecided; {}",
                        state::RUNBOOK
                    ));
                }
                None
            }
        }
    }
}

/// How far a log's turns may be decided for a report of turn `reported`: up to the smaller of
/// `reported` and the log's turns, but for the log's last turn when it has no answer yet. The
/// limit, and whether that last turn was held back.
pub fn decidable(records: &[LogRecord], reported: u64) -> (u64, bool) {
    let count = turns(records);
    let limit = reported.min(count);
    if limit > 0 && limit == count && !turn_records(records, limit).is_some_and(finished) {
        (limit - 1, true)
    } else {
        (limit, false)
    }
}