//! A session on disk: `sessions/<id>/0.baseline.json` and `0.jsonl`. The log is the state —
//! create, open, append and resume all live here, and nothing is durable until `append` syncs it.

use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};

use crate::baseline::{Baseline, BaselineError};
use proto::{CallId, Epoch, LogRecord, SessionId, Timestamp};

/// Creating, opening, appending and resuming a session, and the errors each can raise.
#[derive(Debug)]
pub enum SessionError {
    Exists(SessionId),
    NotFound(SessionId),
    Io(PathBuf, std::io::Error),
    Torn {
        path: PathBuf,
        line: usize,
        why: String,
    },
    Baseline(BaselineError),
    Encode(serde_json::Error),
}

impl std::fmt::Display for SessionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SessionError::Exists(id) => write!(f, "session already exists: {}", id.as_str()),
            SessionError::NotFound(id) => write!(f, "session not found: {}", id.as_str()),
            SessionError::Io(path, err) => write!(f, "{}: {err}", path.display()),
            SessionError::Torn { path, line, why } => write!(
                f,
                "{}:{line}: {why}; see docs/runbook.md#session-log-damaged",
                path.display()
            ),
            SessionError::Baseline(err) => write!(f, "{err}"),
            SessionError::Encode(err) => write!(f, "{err}"),
        }
    }
}

impl std::error::Error for SessionError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            SessionError::Io(_, err) => Some(err),
            SessionError::Baseline(err) => Some(err),
            SessionError::Encode(err) => Some(err),
            _ => None,
        }
    }
}

/// A live session: its id, the directory holding the baseline and log, the replayed records, the
/// appended log file, and the next call id.
#[derive(Debug)]
pub struct Session {
    id: SessionId,
    dir: PathBuf,
    baseline: Baseline,
    records: Vec<LogRecord>,
    log: File,
    next_call: CallId,
}

impl Session {
    pub fn dir_for(home: &Path, id: &SessionId) -> PathBuf {
        home.join("sessions").join(id.as_str())
    }

    pub fn create(
        home: &Path,
        id: SessionId,
        baseline: Baseline,
        slot: u32,
    ) -> Result<Session, SessionError> {
        let dir = Session::dir_for(home, &id);
        if dir.exists() {
            return Err(SessionError::Exists(id));
        }
        std::fs::create_dir_all(&dir).map_err(|e| SessionError::Io(dir.clone(), e))?;

        let baseline_path = dir.join("0.baseline.json");
        let json = baseline.to_json().map_err(SessionError::Encode)?;
        std::fs::write(&baseline_path, json)
            .map_err(|e| SessionError::Io(baseline_path.clone(), e))?;

        let log_path = dir.join("0.jsonl");
        let log = OpenOptions::new()
            .create_new(true)
            .append(true)
            .open(&log_path)
            .map_err(|e| SessionError::Io(log_path.clone(), e))?;

        let mut session = Session {
            id,
            dir,
            baseline,
            records: Vec::new(),
            log,
            next_call: CallId(1),
        };
        session.persist(LogRecord::SessionStart {
            time: Timestamp::now(),
            session: session.id.clone(),
            epoch: Epoch(0),
            slot,
            baseline: session.baseline.hash().map_err(SessionError::Baseline)?,
        })?;
        Ok(session)
    }

    pub fn open(home: &Path, id: SessionId) -> Result<Session, SessionError> {
        let dir = Session::dir_for(home, &id);
        let log_path = dir.join("0.jsonl");
        if !log_path.exists() {
            return Err(SessionError::NotFound(id));
        }
        let baseline_path = dir.join("0.baseline.json");
        let baseline = Baseline::load(&baseline_path).map_err(SessionError::Baseline)?;

        let text = std::fs::read_to_string(&log_path)
            .map_err(|e| SessionError::Io(log_path.clone(), e))?;
        let mut records = Vec::new();
        if !text.is_empty() {
            for (idx, line) in text.split_inclusive('\n').enumerate() {
                let line_no = idx + 1;
                if !line.ends_with('\n') {
                    return Err(SessionError::Torn {
                        path: log_path.clone(),
                        line: line_no,
                        why: "last line does not end in a newline".to_string(),
                    });
                }
                match serde_json::from_str::<LogRecord>(line) {
                    Ok(record) => records.push(record),
                    Err(why) => {
                        return Err(SessionError::Torn {
                            path: log_path.clone(),
                            line: line_no,
                            why: why.to_string(),
                        });
                    }
                }
            }
        }

        let mut highest = 0u64;
        for record in &records {
            if let LogRecord::ToolResult { call, .. } = record {
                highest = highest.max(call.0);
            }
        }

        let log = OpenOptions::new()
            .append(true)
            .open(&log_path)
            .map_err(|e| SessionError::Io(log_path.clone(), e))?;

        Ok(Session {
            id,
            dir,
            baseline,
            records,
            log,
            next_call: CallId(highest + 1),
        })
    }

    pub fn id(&self) -> &SessionId {
        &self.id
    }

    pub fn dir(&self) -> &Path {
        &self.dir
    }

    pub fn baseline(&self) -> &Baseline {
        &self.baseline
    }

    pub fn records(&self) -> &[LogRecord] {
        &self.records
    }

    pub fn append(&mut self, record: LogRecord) -> Result<(), SessionError> {
        self.persist(record)
    }

    pub fn next_call(&mut self) -> CallId {
        let id = self.next_call;
        self.next_call = CallId(self.next_call.0 + 1);
        id
    }

    pub fn last_usage(&self) -> Option<proto::Usage> {
        for record in self.records.iter().rev() {
            if let LogRecord::Usage {
                cache_n,
                prompt_n,
                predicted_n,
                reasoning_tokens,
                thinking_capped,
                ..
            } = record
            {
                return Some(proto::Usage {
                    cache_n: *cache_n,
                    prompt_n: *prompt_n,
                    predicted_n: *predicted_n,
                    reasoning_tokens: *reasoning_tokens,
                    thinking_capped: *thinking_capped,
                });
            }
        }
        None
    }

    /// Encode one record, write it and its newline, sync the data to disk, and only then keep it in
    /// memory — so a record is never in memory without being durable.
    fn persist(&mut self, record: LogRecord) -> Result<(), SessionError> {
        let line = serde_json::to_string(&record).map_err(SessionError::Encode)?;
        let path = self.log_path();
        self.log
            .write_all(line.as_bytes())
            .map_err(|e| SessionError::Io(path.clone(), e))?;
        self.log
            .write_all(b"\n")
            .map_err(|e| SessionError::Io(path.clone(), e))?;
        self.log
            .sync_data()
            .map_err(|e| SessionError::Io(path.clone(), e))?;
        self.records.push(record);
        Ok(())
    }

    fn log_path(&self) -> PathBuf {
        self.dir.join("0.jsonl")
    }
}