use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use serde_json::json;

use super::driver::{AttachmentRef, Driver, Event, EventSink, SessionStatus};
use super::process;
use super::transport::{Launch, Streams, Transport};
use crate::config::{ProviderConfig, SessionConfig};

/// How long to wait for a model to load before giving up. Loading is mostly
/// disk, and a large quantised model on a cold cache is genuinely slow, so this
/// is generous -- the failure it exists for is a server that will never answer.
const READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Message {
    role: String,
    content: String,
}

pub struct LlamaDriver {
    sink: EventSink,
    endpoint: String,
    transcript: PathBuf,
    sampling: serde_json::Map<String, serde_json::Value>,
    cancel: Arc<AtomicBool>,
    session_dir: PathBuf,
}

impl LlamaDriver {
    /// Takes charge of this session's `llama-server`: the one already loaded if
    /// there is one, otherwise a new one.
    #[allow(clippy::too_many_arguments)]
    pub fn launch(
        meta: &SessionConfig,
        provider: &ProviderConfig,
        transport: &Transport,
        models_dir: &Path,
        transcript: &Path,
        session_dir: &Path,
        sink: EventSink,
        _subagents: Arc<super::subagent::Subagents>,
    ) -> Result<Self> {
        let model = meta.model.as_deref().context(
            "a llama.cpp session needs a model -- one of the downloaded ones, by its key",
        )?;
        let path = model_on(transport, models_dir, model)?;

        // Already loaded and still running: keep talking to it. The health poll
        // below confirms it is really answering, so adopting a pid whose server
        // has wedged still reports as a failure rather than as a session that
        // silently never replies.
        if let Some(process::Record {
            detail: process::Detail::Http { port },
            pid,
            ..
        }) = process::live(session_dir)
        {
            tracing::info!(
                "session {} reattaching to the llama-server it left loaded (pid {pid}, port {port})",
                meta.id
            );
            return Ok(Self::attached(
                format!("http://127.0.0.1:{port}"),
                meta,
                model,
                transcript,
                session_dir,
                sink,
            ));
        }

        let forward = transport
            .reserve_port()
            .context("finding a port for llama-server")?;
        let mut args: Vec<String> = vec![
            "-m".into(),
            path.clone(),
            "--host".into(),
            "127.0.0.1".into(),
            "--port".into(),
            forward.there.to_string(),
        ];
        // Settings that belong to the server because they decide how the model
        // is loaded; the sampling ones ride on each request instead, so changing
        // them later needn't reload anything.
        for (key, flag) in [
            ("contextSize", "-c"),
            ("gpuLayers", "-ngl"),
            ("threads", "-t"),
        ] {
            if let Some(value) = meta.params.get(key) {
                args.push(flag.to_string());
                args.push(value.clone());
            }
        }

        let program = provider.program();
        let launch = Launch::new(program, args, meta.cwd.as_deref()).reaching(forward);
        let child = transport.spawn(
            &launch,
            Streams::Detached {
                stdin: std::process::Stdio::null(),
                stdout: log_file(&session_dir.join(SERVER_LOG))?.into(),
                stderr: log_file(&session_dir.join(SERVER_LOG))?.into(),
            },
        )?;
        let pid = child
            .id()
            .context("llama-server exited before it could be recorded")?;
        tracing::info!(
            "session {} running {program} for {model} {} on 127.0.0.1:{} there, \
             reached at 127.0.0.1:{} here, as pid {pid}",
            meta.id,
            transport.describe(),
            forward.there,
            forward.here,
        );
        // Reaped so it does not become a zombie while this server is still its
        // parent; the health poll and the record are what say whether the
        // session is alive, because after a restart there is no `Child` to ask.
        tokio::spawn(async move {
            let mut child = child;
            let _ = child.wait().await;
        });

        // The *near* port, because that is the one anything reaching this
        // server has to dial -- including a later run of this backend,
        // which adopts the record without knowing which machine the server
        // is on. For a remote session the recorded pid is the ssh
        // client's, which is the process this machine owns and which holds
        // the tunnel open for exactly as long as the far server lives.
        let record = process::Record::of(pid, process::Detail::Http { port: forward.here })
            .context("llama-server was gone before its start time could be read")?;
        process::write(session_dir, &record);

        Ok(Self::attached(
            format!("http://127.0.0.1:{}", forward.here),
            meta,
            model,
            transcript,
            session_dir,
            sink,
        ))
    }

    /// Shared by starting one and adopting one, because everything after "there
    /// is a server at this address" is identical -- including waiting for it to
    /// answer, which an adopted one still owes: a recorded pid says a process
    /// exists, not that its model is loaded.
    fn attached(
        endpoint: String,
        meta: &SessionConfig,
        model: &str,
        transcript: &Path,
        session_dir: &Path,
        sink: EventSink,
    ) -> Self {
        let _ = sink.send(Event::Status {
            state: SessionStatus::Running,
        });
        {
            let sink = sink.clone();
            let endpoint = endpoint.clone();
            let model = model.to_string();
            let session_dir = session_dir.to_path_buf();
            std::thread::spawn(move || match wait_until_ready(&endpoint, &session_dir) {
                Ok(()) => {
                    tracing::info!("{model} loaded and answering at {endpoint}");
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Idle,
                    });
                    watch(session_dir, sink);
                }
                Err(err) => {
                    let _ = sink.send(Event::Error {
                        message: format!("{model} never became ready: {err:#}"),
                    });
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Exited,
                    });
                    process::clear(&session_dir);
                }
            });
        }

        let mut sampling = serde_json::Map::new();
        for (key, field) in [
            ("temperature", "temperature"),
            ("topP", "top_p"),
            ("topK", "top_k"),
            ("maxTokens", "max_tokens"),
        ] {
            if let Some(raw) = meta.params.get(key)
                && let Ok(number) = raw.parse::<f64>()
            {
                sampling.insert(field.to_string(), json!(number));
            }
        }

        Self {
            sink,
            endpoint,
            transcript: transcript.to_path_buf(),
            sampling,
            cancel: Arc::new(AtomicBool::new(false)),
            session_dir: session_dir.to_path_buf(),
        }
    }
}

const SERVER_LOG: &str = "llama-server.log";

/// How often a loaded server is checked for still being there. Slower than the
/// Claude driver's stdout poll because nothing is waiting on it: this only has
/// to notice a server that has gone.
const WATCH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(2);

fn log_file(path: &Path) -> Result<std::fs::File> {
    use std::os::unix::fs::OpenOptionsExt;
    std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .mode(0o600)
        .open(path)
        .with_context(|| format!("opening {}", path.display()))
}

fn watch(session_dir: PathBuf, sink: EventSink) {
    std::thread::spawn(move || {
        loop {
            std::thread::sleep(WATCH_INTERVAL);
            match process::recorded(&session_dir) {
                Some((_, process::Liveness::Alive)) => {}
                None => return,
                Some((_, process::Liveness::Dead)) => {
                    let _ = sink.send(Event::Error {
                        message: "llama-server exited".to_string(),
                    });
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Exited,
                    });
                    process::clear(&session_dir);
                    return;
                }
                Some((_, process::Liveness::Unknown)) => {
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Unknown,
                    });
                }
            }
            if sink.is_closed() {
                return;
            }
        }
    });
}

impl Driver for LlamaDriver {
    fn send_user_message(&self, text: String, attachments: Vec<AttachmentRef>) {
        if !attachments.is_empty() {
            let _ = self.sink.send(Event::Error {
                message: "this model can't be sent attachments or files".to_string(),
            });
        }
        let sink = self.sink.clone();
        let endpoint = self.endpoint.clone();
        let transcript = self.transcript.clone();
        let sampling = self.sampling.clone();
        let cancel = Arc::clone(&self.cancel);
        cancel.store(false, Ordering::Relaxed);

        std::thread::spawn(move || {
            // Nothing is ever held back here -- there is no queue to wait in --
            // so the message is taken the moment it arrives. Said anyway,
            // because this is what records it: see `MessageTaken`.
            let _ = sink.send(Event::MessageTaken {
                id: None,
                text: text.clone(),
                // Never any: this driver refuses attachments above.
                attachments: Vec::new(),
            });
            let _ = sink.send(Event::Status {
                state: SessionStatus::Running,
            });
            // Everything before this message, plus this message. Read rather
            // than remembered, and `text` is appended here rather than waited
            // for, because the message's own transcript entry is still on its
            // way when this runs.
            let mut messages = conversation(&transcript);
            messages.push(Message {
                role: "user".into(),
                content: text,
            });
            if let Err(err) = generate(&endpoint, &messages, &sampling, &cancel, &sink) {
                let _ = sink.send(Event::Error {
                    message: format!("{err:#}"),
                });
            }
            let _ = sink.send(Event::Status {
                state: SessionStatus::Idle,
            });
        });
    }

    fn answer_question(&self, _id: &str, _answers: &[String]) {}

    fn interrupt(&self) {
        self.cancel.store(true, Ordering::Relaxed);
    }

    fn set_title(&self, _title: &str) {}

    fn set_permission_mode(&self, _mode: &str) {
        let _ = self.sink.send(Event::Error {
            message: "a llama.cpp session runs no tools, so there is nothing for a permission \
                      mode to govern."
                .to_string(),
        });
    }

    fn set_model(&self, _model: &str) {
        let _ = self.sink.send(Event::Error {
            message: "a llama.cpp session's model is fixed when it starts, because the server \
                      loads one model into memory. Spawn another session to use a different one."
                .to_string(),
        });
    }

    fn run_command(&self, text: &str) {
        let _ = self.sink.send(Event::Error {
            message: format!(
                "a llama.cpp session has no commands of its own, so {text} means nothing to it."
            ),
        });
    }

    fn compact(&self) {
        let _ = self.sink.send(Event::Error {
            message: "llama.cpp has no compaction. Clear the session instead, which costs nothing."
                .to_string(),
        });
    }

    fn clear(&self) {
        let _ = self.sink.send(Event::Cleared);
    }

    /// Worth being deliberate about, because the cost points the other way from
    /// the Claude driver's: a `llama-server` holds its whole model in memory, so
    /// a leaked one is gigabytes nobody is using. It is left anyway, because the
    /// alternative is unloading and reloading that model on every backend
    /// restart -- minutes of disk, for a session somebody is in the middle of.
    /// The record is what keeps it from being *nobody's*.
    fn detach(&self) {
        self.cancel.store(true, Ordering::Relaxed);
    }

    fn stop(&self) {
        self.cancel.store(true, Ordering::Relaxed);
        if let Some(record) = process::live(&self.session_dir) {
            process::stop(&record, process::STOP_GRACE);
        }
        process::clear(&self.session_dir);
    }
}

/// This must stay a pure function of the transcript and must never re-render
/// earlier turns. llama.cpp caches the prompt prefix, so a growing conversation
/// reprocesses almost nothing -- but only while every turn is byte-identical to
/// last time. Changing how an old turn is rendered silently reprocesses the
/// whole history on every message.
fn conversation(path: &Path) -> Vec<Message> {
    let Ok(events) = crate::session::transcript::read_after(path, 0) else {
        return Vec::new();
    };
    let mut messages: Vec<Message> = Vec::new();
    let mut pending = String::new();
    let events = match events.iter().rposition(|e| e.event == Event::Cleared) {
        Some(at) => &events[at + 1..],
        None => &events[..],
    };
    for event in events.iter().cloned() {
        match event.event {
            Event::UserMessage { text, .. } => {
                if !pending.is_empty() {
                    messages.push(Message {
                        role: "assistant".into(),
                        content: std::mem::take(&mut pending),
                    });
                }
                messages.push(Message {
                    role: "user".into(),
                    content: text,
                });
            }
            Event::AssistantText { delta } => pending.push_str(&delta),
            _ => {}
        }
    }
    if !pending.is_empty() {
        messages.push(Message {
            role: "assistant".into(),
            content: pending,
        });
    }
    messages
}

fn model_path(models_dir: &Path, key: &str) -> Result<PathBuf> {
    let mut path = models_dir.to_path_buf();
    for part in key.split('/') {
        if part.is_empty() || part == "." || part == ".." {
            bail!("\"{key}\" is not a model key this can resolve");
        }
        path.push(part);
    }
    if !path.is_file() {
        bail!("no downloaded model called \"{key}\" -- download it first");
    }
    Ok(path)
}

/// One function rather than a local check and hope for the other case: the same
/// question has to be asked of two filesystems. The remote answer is measured
/// for the reason the local one is -- a missing file otherwise becomes a
/// `llama-server` that starts, fails to load, and reports as a session that
/// never became ready, which reads as the machine being slow.
fn model_on(transport: &Transport, models_dir: &Path, key: &str) -> Result<String> {
    let Transport::Ssh { name, .. } = transport else {
        return Ok(model_path(models_dir, key)?.to_string_lossy().into_owned());
    };
    let dir = crate::models::dir_on(transport, models_dir);
    for part in key.split('/') {
        if part.is_empty() || part == "." || part == ".." {
            bail!("\"{key}\" is not a model key this can resolve");
        }
    }
    let path = format!("{}/{key}", dir.trim_end_matches('/'));
    // `$HOME` on the far side, which is the only machine that knows what it is,
    // and the resolved path printed back so the launch hands `llama-server`
    // something absolute. "Not there" is answered rather than failed, because a
    // machine that could not be asked at all has to say so in its own words --
    // it would otherwise arrive as this same sentence about a missing model.
    let script = "p=$1; case $p in \"~\") p=$HOME;; \"~/\"*) p=$HOME/${p#\"~/\"};; esac; \
                  [ -f \"$p\" ] && printf 'at\\t%s\\n' \"$p\" || printf 'missing\\n'"
        .to_string();
    let launch = Launch::new(
        "sh",
        vec!["-c".to_string(), script, "sh".to_string(), path.clone()],
        None,
    );
    let answer = transport
        .capture_blocking(&launch)
        .with_context(|| format!("couldn't ask {name} where its models are"))?;
    match answer.trim().split_once('\t') {
        Some(("at", resolved)) => Ok(resolved.to_string()),
        _ => bail!(
            "{name} has no model at {path}. A llama.cpp session serves the file from the \
             machine it runs on, so the model has to be on {name} -- what this backend has \
             downloaded is somewhere else."
        ),
    }
}

/// Watches the process as well as the port, because the two failures need
/// different words and one of them is common: a model that will not load,
/// a port already taken on the far machine, a `llama-server` too old for
/// a flag. All of those exit within a second and none of them will ever
/// answer `/health`, so waiting out the timeout turns a server that said
/// exactly what was wrong into "gave up after 300s".
fn wait_until_ready(endpoint: &str, session_dir: &Path) -> Result<()> {
    let deadline = std::time::Instant::now() + READY_TIMEOUT;
    let url = format!("{endpoint}/health");
    loop {
        if let Ok(response) = ureq::get(&url).call()
            && response.status() == 200
        {
            return Ok(());
        }
        // `None` is the session having been stopped or deleted while this
        // waited, which is nobody's fault and still not worth waiting on.
        match process::recorded(session_dir) {
            Some((_, process::Liveness::Alive | process::Liveness::Unknown)) => {}
            Some((_, process::Liveness::Dead)) | None => {
                bail!("it exited before it answered.{}", log_tail(session_dir));
            }
        }
        if std::time::Instant::now() > deadline {
            bail!(
                "gave up after {}s.{}",
                READY_TIMEOUT.as_secs(),
                log_tail(session_dir)
            );
        }
        std::thread::sleep(std::time::Duration::from_millis(250));
    }
}

/// Its account of what went wrong is the useful half -- "failed to load
/// model", "bind: Address already in use" -- and on a remote session it
/// is the only half, since nobody reading the phone can open a file on
/// that machine. Bounded, because this ends up in an event a phone draws.
fn log_tail(session_dir: &Path) -> String {
    let Ok(text) = std::fs::read_to_string(session_dir.join(SERVER_LOG)) else {
        return String::new();
    };
    let tail: Vec<&str> = text.lines().rev().take(LOG_TAIL_LINES).collect();
    if tail.is_empty() {
        return String::new();
    }
    format!(
        " It last said: {}",
        tail.into_iter().rev().collect::<Vec<_>>().join(" / ")
    )
}

const LOG_TAIL_LINES: usize = 6;

/// One streamed completion: posts the conversation, emits each delta as it
/// arrives. Emits rather than returns, because the transcript those events land
/// in is what the next turn reads back.
fn generate(
    endpoint: &str,
    messages: &[Message],
    sampling: &serde_json::Map<String, serde_json::Value>,
    cancel: &AtomicBool,
    sink: &EventSink,
) -> Result<()> {
    let mut body = json!({
        "messages": messages,
        "stream": true,
        "stream_options": {"include_usage": true},
    });
    let map = body.as_object_mut().expect("built as an object");
    for (key, value) in sampling {
        map.insert(key.clone(), value.clone());
    }

    let mut response = ureq::post(format!("{endpoint}/v1/chat/completions"))
        .header("Content-Type", "application/json")
        .send_json(&body)
        .context("asking llama-server to generate")?;

    let reader = std::io::BufReader::new(response.body_mut().as_reader());
    let mut tokens = 0u64;
    let mut context = None;
    for line in std::io::BufRead::lines(reader) {
        if cancel.load(Ordering::Relaxed) {
            break;
        }
        let line = line.context("reading the generation stream")?;
        let Some(payload) = line.strip_prefix("data: ") else {
            continue;
        };
        if payload.trim() == "[DONE]" {
            break;
        }
        let Ok(chunk) = serde_json::from_str::<serde_json::Value>(payload) else {
            continue;
        };
        if let Some(usage) = chunk.get("usage") {
            if let Some(total) = usage
                .get("total_tokens")
                .and_then(serde_json::Value::as_u64)
            {
                tokens = total;
            }
            if let Some(prompt) = usage
                .get("prompt_tokens")
                .and_then(serde_json::Value::as_u64)
            {
                context = Some(prompt);
            }
        }
        let delta = chunk
            .get("choices")
            .and_then(|c| c.get(0))
            .and_then(|c| c.get("delta"))
            .and_then(|d| d.get("content"))
            .and_then(serde_json::Value::as_str)
            .unwrap_or_default();
        if !delta.is_empty() {
            let _ = sink.send(Event::AssistantText {
                delta: delta.to_string(),
            });
        }
    }
    if tokens > 0 {
        let _ = sink.send(Event::UsageDelta { tokens, context });
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::session::transcript::Transcript;

    fn transcript_with(events: &[Event]) -> (tempfile::TempDir, PathBuf) {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("transcript.jsonl");
        let mut transcript = Transcript::open(&path).expect("open");
        for event in events {
            transcript.append(event.clone(), 0.0).expect("append");
        }
        (dir, path)
    }

    #[test]
    fn deltas_between_user_messages_are_one_assistant_turn() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "hello".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "hi ".into(),
            },
            Event::AssistantText {
                delta: "there".into(),
            },
            Event::Status {
                state: SessionStatus::Idle,
            },
            Event::UserMessage {
                id: None,
                text: "again".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "yes".into(),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(
            messages
                .iter()
                .map(|m| (m.role.as_str(), m.content.as_str()))
                .collect::<Vec<_>>(),
            [
                ("user", "hello"),
                ("assistant", "hi there"),
                ("user", "again"),
                ("assistant", "yes")
            ],
        );
    }

    #[test]
    fn an_interrupted_reply_stays_in_the_conversation() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "count".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "one two".into(),
            },
            Event::Status {
                state: SessionStatus::Idle,
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[1].content, "one two");
    }

    #[test]
    fn other_events_are_not_part_of_the_conversation() {
        let (_dir, path) = transcript_with(&[
            Event::Status {
                state: SessionStatus::Running,
            },
            Event::UserMessage {
                id: None,
                text: "hello".into(),
                attachments: Vec::new(),
            },
            Event::Error {
                message: "something went wrong".into(),
            },
            Event::AssistantText {
                delta: "still here".into(),
            },
            Event::UsageDelta {
                tokens: 12,
                context: Some(12),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].content, "hello");
        assert_eq!(messages[1].content, "still here");
    }

    #[test]
    fn the_conversation_starts_after_the_last_clear() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "the long expensive conversation".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "at length".into(),
            },
            Event::Cleared,
            Event::UserMessage {
                id: None,
                text: "a fresh start".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "cheaply".into(),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].content, "a fresh start");
        assert_eq!(messages[1].content, "cheaply");
    }

    #[test]
    fn only_the_newest_clear_counts() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "one".into(),
                attachments: Vec::new(),
            },
            Event::Cleared,
            Event::UserMessage {
                id: None,
                text: "two".into(),
                attachments: Vec::new(),
            },
            Event::Cleared,
            Event::UserMessage {
                id: None,
                text: "three".into(),
                attachments: Vec::new(),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 1);
        assert_eq!(messages[0].content, "three");
    }

    #[test]
    fn a_model_key_cannot_climb_out_of_the_models_directory() {
        let dir = tempfile::tempdir().expect("tempdir");
        for attempt in ["../../etc/passwd", "unsloth/../../escape.gguf", ""] {
            assert!(
                model_path(dir.path(), attempt).is_err(),
                "{attempt:?} should have been refused",
            );
        }
    }
}