use std::io::{self, IsTerminal, Read, Write};
use std::path::PathBuf;

use std::sync::{Arc, mpsc};
use std::thread;

use xserv_model::{
    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, SamplingParams,
    loader, sample, sample_greedy_penalized,
};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;

enum ChatModel {
    Qwen3(Qwen3),
    GptOss(GptOss),
}

impl ChatModel {
    fn forward_prefill_paged(
        &self,
        tokens: &[u32],
        slot: usize,
        cache: &mut PagedKVCache,
    ) -> xserv_tensor::Tensor {
        match self {
            ChatModel::Qwen3(m) => m.forward_prefill_paged(tokens, slot, cache),
            ChatModel::GptOss(m) => m.forward_prefill_paged(tokens, slot, cache),
        }
    }
    fn forward_decode_paged(
        &self,
        tokens: &[u32],
        positions: &[usize],
        slots: &[usize],
        cache: &mut PagedKVCache,
    ) -> xserv_tensor::Tensor {
        match self {
            ChatModel::Qwen3(m) => m.forward_decode_paged(tokens, positions, slots, cache),
            ChatModel::GptOss(m) => m.forward_decode_paged(tokens, positions, slots, cache),
        }
    }
}

// TP worker infrastructure (reused from tp_engine pattern)
#[derive(Clone)]
enum TpCommand {
    Register(usize),
    Free(usize),
    Prefill {
        tokens: Vec<u32>,
        slot: usize,
    },
    Decode {
        tokens: Vec<u32>,
        positions: Vec<usize>,
        slots: Vec<usize>,
    },
}

struct TpHandle {
    cmd_txs: Vec<mpsc::Sender<TpCommand>>,
    ack_rx: mpsc::Receiver<()>,
}

impl TpHandle {
    fn send(&self, cmd: TpCommand) {
        for tx in &self.cmd_txs {
            tx.send(cmd.clone()).ok();
        }
    }
    fn wait(&self) {
        for _ in 0..self.cmd_txs.len() {
            self.ack_rx.recv().ok();
        }
    }
}

fn tp_worker_loop(
    rank: usize,
    world: usize,
    id: xserv_distributed::UniqueId,
    model_dir: std::path::PathBuf,
    config: ModelConfig,
    max_seq_len: usize,
    cmd_rx: mpsc::Receiver<TpCommand>,
    ack_tx: mpsc::Sender<()>,
) {
    let tp = Arc::new(xserv_distributed::TpContext::init(
        rank,
        world,
        id,
        rank as u32,
    ));
    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
    let model = if config.is_moe() {
        ChatModel::GptOss(GptOss::from_weights_tp(
            config.clone(),
            weights,
            rank,
            world,
            rank as u32,
            Some(tp),
        ))
    } else {
        ChatModel::Qwen3(Qwen3::from_weights_tp(
            config.clone(),
            weights,
            rank,
            world,
            rank as u32,
            Some(tp),
        ))
    };
    let local_kv = config.num_kv_heads() / world;
    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
    let total_blocks = max_blocks_per_seq + 8;
    let mut cache = PagedKVCache::new_tp(
        &config,
        local_kv,
        total_blocks,
        0,
        1,
        max_blocks_per_seq,
        DType::BF16,
        rank as u32,
    );
    let mut decoder = GraphedGptOssDecoder::new();
    while let Ok(cmd) = cmd_rx.recv() {
        match cmd {
            TpCommand::Register(slot) => {
                let _ = cache.register_sequence(slot);
            }
            TpCommand::Free(slot) => cache.free_sequence(slot),
            TpCommand::Prefill { tokens, slot } => {
                let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
            }
            TpCommand::Decode {
                tokens,
                positions,
                slots,
            } => {
                let _ = chat_decode(
                    &model,
                    &mut decoder,
                    &tokens,
                    &positions,
                    &slots,
                    &mut cache,
                );
            }
        }
        let _ = ack_tx.send(());
    }
}

const SLOT: usize = 0;

struct CliOptions {
    model_dir: PathBuf,
    max_tokens: usize,
    max_seq_len: usize,
    tp: usize,
    sampling: SamplingParams,
    system_prompt: Option<String>,
    enable_thinking: bool,
    color: bool,
}

enum Finish {
    Stop { token_id: u32 },
    Length,
}

enum Line {
    Text(String),
    Eof,
}

/// RAII terminal raw-mode guard. Disables canonical mode + echo (keeps output
/// post-processing and signals), so we read keystrokes ourselves and edit the
/// line UTF-8-aware. Restores the original termios on drop.
struct RawMode {
    orig: libc::termios,
}

impl RawMode {
    fn enable() -> Option<Self> {
        unsafe {
            let mut orig: libc::termios = std::mem::zeroed();
            if libc::tcgetattr(libc::STDIN_FILENO, &mut orig) != 0 {
                return None;
            }
            let mut raw = orig;
            raw.c_lflag &= !(libc::ICANON | libc::ECHO);
            raw.c_cc[libc::VMIN as usize] = 1;
            raw.c_cc[libc::VTIME as usize] = 0;
            if libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &raw) != 0 {
                return None;
            }
            Some(RawMode { orig })
        }
    }
}

impl Drop for RawMode {
    fn drop(&mut self) {
        unsafe {
            libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &self.orig);
        }
    }
}

/// Read one line with UTF-8/CJK-aware editing. In a TTY this enters raw mode and
/// handles keystrokes so Backspace deletes a whole character (not a byte), and
/// multi-byte input (汉字/日本語/한글) renders correctly. Non-TTY (piped) input
/// falls back to a plain cooked read.
fn read_line_edited(prompt: &str) -> Line {
    let cooked = || -> Line {
        print!("{prompt}");
        io::stdout().flush().ok();
        let mut s = String::new();
        match io::stdin().read_line(&mut s) {
            Ok(0) | Err(_) => Line::Eof,
            Ok(_) => Line::Text(s),
        }
    };

    if !io::stdin().is_terminal() {
        return cooked();
    }
    let Some(_raw) = RawMode::enable() else {
        return cooked();
    };

    // Single-line editor: on every edit, rewrite the whole line so the terminal
    // renders correct (incl. double-width CJK) glyphs; \x1b[K clears leftovers.
    let redraw = |buf: &str| {
        print!("\r{prompt}{buf}\x1b[K");
        io::stdout().flush().ok();
    };

    let mut buf = String::new();
    redraw(&buf);
    let mut stdin = io::stdin().lock();
    let mut byte = [0u8; 1];

    loop {
        if stdin.read(&mut byte).unwrap_or(0) == 0 {
            // EOF on the stream.
            if buf.is_empty() {
                return Line::Eof;
            }
            break;
        }
        match byte[0] {
            b'\r' | b'\n' => {
                println!();
                break;
            }
            0x7f | 0x08 => {
                // Backspace: drop one whole char (String::pop is char-aware).
                buf.pop();
                redraw(&buf);
            }
            0x04 => {
                // Ctrl-D: EOF only when the line is empty.
                if buf.is_empty() {
                    return Line::Eof;
                }
            }
            0x1b => {
                // Escape sequence (arrows, etc.): consume and ignore the 2 bytes
                // of a typical CSI sequence so they don't land in the buffer.
                let mut seq = [0u8; 2];
                let _ = stdin.read(&mut seq);
            }
            b if b < 0x20 => { /* other control bytes: ignore */ }
            b if b < 0x80 => {
                buf.push(b as char);
                redraw(&buf);
            }
            b => {
                // UTF-8 multi-byte: read the continuation bytes for this char.
                let extra = if b >= 0xF0 {
                    3
                } else if b >= 0xE0 {
                    2
                } else {
                    1
                };
                let mut bytes = vec![b];
                let mut cont = [0u8; 1];
                let mut ok = true;
                for _ in 0..extra {
                    if stdin.read(&mut cont).unwrap_or(0) == 0 {
                        ok = false;
                        break;
                    }
                    bytes.push(cont[0]);
                }
                if ok {
                    if let Ok(s) = std::str::from_utf8(&bytes) {
                        buf.push_str(s);
                        redraw(&buf);
                    }
                }
            }
        }
    }
    Line::Text(buf)
}

fn main() {
    let opts = parse_args();

    xserv_cuda::device::set_device(0).unwrap();
    let info = xserv_cuda::device::device_info(0).unwrap();
    eprintln!(
        "GPU: {} ({} MB free)",
        info.name,
        info.free_memory / 1024 / 1024
    );

    let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
    let model_type = config.model_type.as_deref().unwrap_or("unknown");
    let is_moe = config.is_moe();

    let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
    eprintln!(
        "Model: {model_type}{}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
        if is_moe { " (MoE)" } else { "" },
        config.num_layers(),
        config.hidden(),
        config.num_heads(),
        config.num_kv_heads(),
        config.vocab_size,
        max_seq_len
    );

    let world = opts.tp;
    if world > 1 {
        assert!(
            config.num_kv_heads() % world == 0,
            "num_kv_heads {} not divisible by tp {world}",
            config.num_kv_heads()
        );
    }

    let (model, mut cache, tp_handle) = if world > 1 {
        let id = xserv_distributed::get_unique_id();
        let (ack_tx, ack_rx) = mpsc::channel::<()>();
        let mut cmd_txs = Vec::new();
        for rank in 1..world {
            let (ctx_tx, ctx_rx) = mpsc::channel::<TpCommand>();
            cmd_txs.push(ctx_tx);
            let ack_tx = ack_tx.clone();
            let model_dir = opts.model_dir.clone();
            let config = config.clone();
            thread::spawn(move || {
                tp_worker_loop(
                    rank,
                    world,
                    id,
                    model_dir,
                    config,
                    max_seq_len,
                    ctx_rx,
                    ack_tx,
                );
            });
        }
        eprintln!("Loading weights (tp={world})...");
        let tp = Arc::new(xserv_distributed::TpContext::init(0, world, id, 0));
        let weights = loader::load_model_dir(&opts.model_dir, Device::Cpu);
        eprintln!("Loaded {} tensors", weights.len());
        let m = if is_moe {
            ChatModel::GptOss(GptOss::from_weights_tp(
                config.clone(),
                weights,
                0,
                world,
                0,
                Some(tp),
            ))
        } else {
            ChatModel::Qwen3(Qwen3::from_weights_tp(
                config.clone(),
                weights,
                0,
                world,
                0,
                Some(tp),
            ))
        };
        let local_kv = config.num_kv_heads() / world;
        let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
        let total_blocks = max_blocks_per_seq + 8;
        let c = PagedKVCache::new_tp(
            &config,
            local_kv,
            total_blocks,
            0,
            1,
            max_blocks_per_seq,
            DType::BF16,
            0,
        );
        let h = TpHandle { cmd_txs, ack_rx };
        (m, c, Some(h))
    } else {
        eprintln!("Loading weights...");
        let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
        eprintln!("Loaded {} tensors", weights.len());
        let m = if is_moe {
            ChatModel::GptOss(GptOss::from_weights(config.clone(), weights))
        } else {
            ChatModel::Qwen3(Qwen3::from_weights(config.clone(), weights))
        };
        let c = new_paged_cache(&config, max_seq_len);
        (m, c, None)
    };

    let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
    let mut decoder = GraphedGptOssDecoder::new();
    if let Some(h) = &tp_handle {
        h.send(TpCommand::Register(SLOT));
        h.wait();
    }
    cache.register_sequence(SLOT).expect("register chat slot");
    let use_color = opts.color && io::stdout().is_terminal();

    eprintln!("Ready (paged KV cache, tp={world}).");
    eprintln!("Commands: /exit, /quit, /clear\n");

    // gpt-oss multi-turn history of (user, assistant-final) text. Harmony
    // requires re-rendering the conversation each turn with prior analysis
    // dropped, so the moe path re-prefills from this rather than reusing an
    // incremental KV cache (which would accumulate CoT + <|return|> and collapse
    // at longer context). Qwen3 ignores this and keeps the incremental cache.
    let mut moe_history: Vec<(String, String)> = Vec::new();

    loop {
        let line = match read_line_edited("user> ") {
            Line::Eof => break,
            Line::Text(s) => s,
        };
        let input = line.trim();
        if input.is_empty() {
            continue;
        }

        match input {
            "/exit" | "/quit" | "exit" | "quit" => break,
            "/clear" => {
                reset_slot(&mut cache, &tp_handle);
                moe_history.clear();
                eprintln!("history and KV cache cleared");
                continue;
            }
            "/help" => {
                print_help();
                continue;
            }
            _ => {}
        }

        if is_moe {
            // Harmony multi-turn: re-render the whole conversation (prior
            // analysis dropped) and re-prefill into a freshly cleared slot.
            let prompt =
                build_conversation_gpt_oss(opts.system_prompt.as_deref(), &moe_history, input);
            let prompt_tokens = tokenizer.encode(&prompt);
            if prompt_tokens.is_empty() {
                continue;
            }
            if prompt_tokens.len() >= max_seq_len {
                eprintln!(
                    "context full: conversation needs {} tokens >= max_seq_len {max_seq_len}; use /clear",
                    prompt_tokens.len()
                );
                continue;
            }
            let max_new_tokens = opts.max_tokens.min(max_seq_len - prompt_tokens.len());
            reset_slot(&mut cache, &tp_handle);
            print!("assistant> ");
            io::stdout().flush().unwrap();
            let (_finish, answer) = generate_with_paged_cache(
                &model,
                &mut decoder,
                &mut cache,
                &tokenizer,
                &prompt_tokens,
                &opts.sampling,
                max_new_tokens,
                use_color,
                &tp_handle,
                is_moe,
                opts.enable_thinking,
            );
            moe_history.push((input.to_string(), answer));
            println!();
            continue;
        }

        // Qwen3: incremental KV cache — only the new turn is prefilled and the
        // assistant's tokens stay cached for the next turn.
        let include_system = cache.seq_len(SLOT) == 0;
        let prompt = build_turn_prompt(
            opts.system_prompt.as_deref(),
            include_system,
            input,
            opts.enable_thinking,
        );
        let prompt_tokens = tokenizer.encode(&prompt);
        if prompt_tokens.is_empty() {
            continue;
        }

        let used = cache.seq_len(SLOT);
        let remaining = max_seq_len.saturating_sub(used);
        if prompt_tokens.len() >= remaining {
            eprintln!(
                "context full: {used}/{max_seq_len} tokens used, new turn needs {} tokens; use /clear",
                prompt_tokens.len()
            );
            continue;
        }
        let max_new_tokens = opts.max_tokens.min(remaining - prompt_tokens.len());

        print!("assistant> ");
        io::stdout().flush().unwrap();
        let (finish, _answer) = generate_with_paged_cache(
            &model,
            &mut decoder,
            &mut cache,
            &tokenizer,
            &prompt_tokens,
            &opts.sampling,
            max_new_tokens,
            use_color,
            &tp_handle,
            is_moe,
            opts.enable_thinking,
        );
        match finish {
            Finish::Stop { token_id } => {
                append_after_stop(
                    &model,
                    &mut cache,
                    &tokenizer,
                    max_seq_len,
                    token_id,
                    &tp_handle,
                );
            }
            Finish::Length => {
                append_text_to_cache(
                    &model,
                    &mut cache,
                    &tokenizer,
                    max_seq_len,
                    "<|im_end|>\n",
                    &tp_handle,
                );
            }
        }
        println!();
    }
}

/// Free and re-register the single chat KV slot (clears all cached context).
fn reset_slot(cache: &mut PagedKVCache, tp: &Option<TpHandle>) {
    if let Some(h) = tp {
        h.send(TpCommand::Free(SLOT));
        h.wait();
    }
    cache.free_sequence(SLOT);
    if let Some(h) = tp {
        h.send(TpCommand::Register(SLOT));
        h.wait();
    }
    cache.register_sequence(SLOT).expect("register chat slot");
}

fn parse_args() -> CliOptions {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
        print_usage_and_exit(0);
    }

    let mut model_dir = None;
    let mut max_tokens = 256usize;
    let mut max_seq_len = 2048usize;
    let mut tp = 1usize;
    let mut temperature = 0.0f32;
    let mut top_k = 0usize;
    let mut top_p = 1.0f32;
    let mut system_prompt = None;
    let mut enable_thinking = false;
    let mut color = true;

    let mut i = 0;
    while i < args.len() {
        match args[i].as_str() {
            "-m" | "--model" => {
                i += 1;
                model_dir = args.get(i).map(PathBuf::from);
            }
            "--max-tokens" => {
                i += 1;
                max_tokens = parse_value(&args, i, "--max-tokens");
            }
            "--max-seq-len" => {
                i += 1;
                max_seq_len = parse_value(&args, i, "--max-seq-len");
            }
            "--tp" => {
                i += 1;
                tp = parse_value(&args, i, "--tp");
            }
            "--temperature" => {
                i += 1;
                temperature = parse_value(&args, i, "--temperature");
            }
            "--top-k" => {
                i += 1;
                top_k = parse_value(&args, i, "--top-k");
            }
            "--top-p" => {
                i += 1;
                top_p = parse_value(&args, i, "--top-p");
            }
            "--system" => {
                i += 1;
                system_prompt = args.get(i).cloned();
                if system_prompt.is_none() {
                    eprintln!("missing value for --system");
                    std::process::exit(2);
                }
            }
            "--think" => {
                enable_thinking = true;
            }
            "--no-color" => {
                color = false;
            }
            arg if arg.starts_with('-') => {
                eprintln!("unknown option: {arg}");
                print_usage_and_exit(2);
            }
            arg => {
                if model_dir.is_some() {
                    eprintln!("unexpected extra argument: {arg}");
                    print_usage_and_exit(2);
                }
                model_dir = Some(PathBuf::from(arg));
            }
        }
        i += 1;
    }

    CliOptions {
        model_dir: model_dir.unwrap_or_else(|| {
            eprintln!("missing model directory");
            print_usage_and_exit(2);
        }),
        max_tokens: max_tokens.max(1),
        max_seq_len: max_seq_len.max(1),
        tp: tp.max(1),
        sampling: SamplingParams {
            temperature,
            top_k,
            top_p,
        },
        system_prompt,
        enable_thinking,
        color,
    }
}

fn parse_value<T: std::str::FromStr>(args: &[String], i: usize, name: &str) -> T {
    args.get(i).and_then(|s| s.parse().ok()).unwrap_or_else(|| {
        eprintln!("invalid or missing value for {name}");
        std::process::exit(2);
    })
}

fn print_usage_and_exit(code: i32) -> ! {
    eprintln!(
        "Usage: xserv-chat <model-dir> [options]\n\
         \n\
         Options:\n\
         \t-m, --model DIR       Model directory\n\
         \t--max-tokens N        Max generated tokens per turn (default: 256)\n\
         \t--max-seq-len N       Persistent KV context length (default: 2048)\n\
         \t--tp N                Tensor parallelism degree (default: 1)\n\
         \t--temperature F       Sampling temperature, 0 = greedy (default: 0)\n\
         \t--top-k N             Top-k sampling, 0 = disabled (default: 0)\n\
         \t--top-p F             Top-p sampling (default: 1.0)\n\
         \t--system TEXT         System prompt for the first turn after start or /clear\n\
         \t--think               Let Qwen3 emit thinking; rendered gray on terminals\n\
         \t--no-color            Disable ANSI color for thinking output\n\
         \t-h, --help            Show this help"
    );
    std::process::exit(code);
}

fn print_help() {
    eprintln!("Commands:");
    eprintln!("  /clear  clear chat history and free/recreate the paged KV slot");
    eprintln!("  /exit   quit");
    eprintln!("  /quit   quit");
}

fn new_paged_cache(config: &ModelConfig, max_seq_len: usize) -> PagedKVCache {
    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
    let total_blocks = (max_blocks_per_seq + 1).max(2);
    // Single-slot interactive CLI: no swap pool (cpu_total_blocks = 0).
    PagedKVCache::new(
        config,
        total_blocks,
        0,
        1,
        max_blocks_per_seq,
        DType::BF16,
        0,
    )
}

fn build_turn_prompt(
    system: Option<&str>,
    include_system: bool,
    user_input: &str,
    enable_thinking: bool,
) -> String {
    let mut prompt = String::new();
    if include_system {
        if let Some(system) = system {
            if !system.trim().is_empty() {
                prompt.push_str("<|im_start|>system\n");
                prompt.push_str(system.trim());
                prompt.push_str("<|im_end|>\n");
            }
        }
    }
    prompt.push_str("<|im_start|>user\n");
    prompt.push_str(user_input);
    prompt.push_str("<|im_end|>\n");
    prompt.push_str("<|im_start|>assistant\n");
    if !enable_thinking {
        prompt.push_str("<think>\n\n</think>\n\n");
    }
    prompt
}

/// Render the full gpt-oss harmony conversation for re-prefill. gpt-oss was
/// trained on this exact system-message structure (identity / knowledge cutoff
/// / current date / Reasoning level / channels — see the model's
/// chat_template.jinja `build_system_message`). A hand-rolled substitute puts
/// the model out of distribution and destabilizes channel selection.
///
/// Harmony multi-turn drops prior chain-of-thought: past assistant messages are
/// rendered as completed `final` channels ending in `<|end|>` (not the
/// `<|return|>` stop token). Keeping the analysis + `<|return|>` of every turn
/// in context — as an incremental KV cache does — is out of distribution and
/// makes the model collapse at longer context. "Reasoning: low" keeps the
/// analysis channel short for an interactive chat.
fn build_conversation_gpt_oss(
    system: Option<&str>,
    history: &[(String, String)],
    current_user: &str,
) -> String {
    let mut prompt = String::new();
    prompt.push_str("<|start|>system<|message|>");
    prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
    prompt.push_str("Knowledge cutoff: 2024-06\n");
    prompt.push_str(&format!("Current date: {}\n\n", today_ymd()));
    prompt.push_str("Reasoning: low\n\n");
    prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
    prompt.push_str("<|end|>");
    if let Some(sys) = system {
        if !sys.trim().is_empty() {
            prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
            prompt.push_str(sys.trim());
            prompt.push_str("<|end|>");
        }
    }
    for (user, assistant) in history {
        prompt.push_str("<|start|>user<|message|>");
        prompt.push_str(user);
        prompt.push_str("<|end|>");
        prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
        prompt.push_str(assistant.trim());
        prompt.push_str("<|end|>");
    }
    prompt.push_str("<|start|>user<|message|>");
    prompt.push_str(current_user);
    prompt.push_str("<|end|>");
    prompt.push_str("<|start|>assistant");
    prompt
}

/// Current UTC date as "YYYY-MM-DD" for the harmony system message. Rata Die
/// civil-calendar conversion (same algorithm the server uses for strftime_now).
fn today_ymd() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs();
    let z = (secs / 86400) as i64 + 719468;
    let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
    let doe = z - era * 146097;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    format!("{y:04}-{m:02}-{d:02}")
}

fn chat_decode(
    model: &ChatModel,
    decoder: &mut GraphedGptOssDecoder,
    tokens: &[u32],
    positions: &[usize],
    slots: &[usize],
    cache: &mut PagedKVCache,
) -> xserv_tensor::Tensor {
    match model {
        ChatModel::GptOss(m) => decoder.decode(m, tokens, positions, slots, cache),
        ChatModel::Qwen3(_) => model.forward_decode_paged(tokens, positions, slots, cache),
    }
}

fn generate_with_paged_cache(
    model: &ChatModel,
    decoder: &mut GraphedGptOssDecoder,
    cache: &mut PagedKVCache,
    tokenizer: &Tokenizer,
    prompt_tokens: &[u32],
    sampling: &SamplingParams,
    max_tokens: usize,
    use_color: bool,
    tp: &Option<TpHandle>,
    is_moe: bool,
    enable_thinking: bool,
) -> (Finish, String) {
    let harmony_end_id = if is_moe {
        tokenizer.special_token_id("<|end|>")
    } else {
        None
    };
    let harmony_channel_id = if is_moe {
        tokenizer.special_token_id("<|channel|>")
    } else {
        None
    };
    let harmony_message_id = if is_moe {
        tokenizer.special_token_id("<|message|>")
    } else {
        None
    };
    let harmony_special: Vec<u32> = if is_moe {
        [
            "<|channel|>",
            "<|start|>",
            "<|end|>",
            "<|message|>",
            "<|return|>",
        ]
        .iter()
        .filter_map(|s| tokenizer.special_token_id(s))
        .collect()
    } else {
        Vec::new()
    };
    // Harmony channel state: "final" channel text is printed normally,
    // "analysis" channel is rendered as thinking (gray). After <|channel|>
    // we read the channel name tokens until <|message|>.
    #[derive(PartialEq, Clone, Copy)]
    enum HarmonyState {
        Normal,
        ReadingChannel,
        InAnalysis,
        InFinal,
    }
    let mut hstate = if is_moe {
        HarmonyState::InFinal
    } else {
        HarmonyState::Normal
    };

    // Off by default. A repetition penalty over a harmony stream penalizes the
    // control tokens (<|channel|>, <|message|>, <|start|>) that MUST repeat to
    // open the final channel — so a non-1.0 default makes gpt-oss stop right
    // after the analysis block, before emitting any answer. Opt in via the env
    // var if you want it for plain (non-harmony) generation.
    let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(1.0);
    let rep_window: usize = std::env::var("XSERV_REP_WINDOW")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(512);
    let mut history: Vec<u32> = Vec::new();

    let pick = |logits: &xserv_tensor::Tensor, sp: &SamplingParams, hist: &[u32]| -> u32 {
        if rep_penalty > 1.0 && sp.temperature == 0.0 {
            let start = hist.len().saturating_sub(rep_window);
            sample_greedy_penalized(logits, &hist[start..], rep_penalty)
        } else {
            sample(logits, sp)
        }
    };

    if let Some(h) = tp {
        h.send(TpCommand::Prefill {
            tokens: prompt_tokens.to_vec(),
            slot: SLOT,
        });
    }
    let logits = model.forward_prefill_paged(prompt_tokens, SLOT, cache);
    if let Some(h) = tp {
        h.wait();
    }
    let mut next = pick(&logits, sampling, &history);
    let mut decode_buffer = Vec::new();
    let mut in_thinking = false;
    let show_thinking = is_moe && enable_thinking;
    // Visible answer tokens, returned for multi-turn history. For moe this is
    // the final-channel content only (analysis is suppressed/gray); for Qwen3
    // it is everything printed. The caller decodes these into the assistant
    // message it re-renders into the next prompt.
    let mut answer_ids: Vec<u32> = Vec::new();

    for _ in 0..max_tokens {
        let position = cache.seq_len(SLOT);
        if let Some(h) = tp {
            h.send(TpCommand::Decode {
                tokens: vec![next],
                positions: vec![position],
                slots: vec![SLOT],
            });
        }
        let logits = chat_decode(model, decoder, &[next], &[position], &[SLOT], cache);
        if let Some(h) = tp {
            h.wait();
        }
        if tokenizer.is_eos(next) {
            print_stream_text(
                &tokenizer.flush_decode_stream(&mut decode_buffer),
                in_thinking,
                use_color,
            );
            if show_thinking && in_thinking {
                print_stream_text("\n</think>\n\n", true, use_color);
            }
            io::stdout().flush().unwrap();
            return (
                Finish::Stop { token_id: next },
                tokenizer.decode(&answer_ids),
            );
        }
        if harmony_end_id == Some(next) {
            // <|end|> closes current segment; if in final channel, we're done
            print_stream_text(
                &tokenizer.flush_decode_stream(&mut decode_buffer),
                in_thinking,
                use_color,
            );
            if hstate == HarmonyState::InFinal {
                io::stdout().flush().unwrap();
                return (
                    Finish::Stop { token_id: next },
                    tokenizer.decode(&answer_ids),
                );
            }
            // Closing a thinking (analysis/commentary) channel: emit the </think>
            // marker so it renders like Qwen3's thinking block.
            if show_thinking && hstate == HarmonyState::InAnalysis {
                print_stream_text("\n</think>\n\n", true, use_color);
                in_thinking = false;
            }
            hstate = HarmonyState::Normal;
            next = pick(&logits, sampling, &history);
            continue;
        }

        history.push(next);

        // Harmony channel routing state machine
        if harmony_channel_id == Some(next) {
            decode_buffer.clear();
            hstate = HarmonyState::ReadingChannel;
            next = pick(&logits, sampling, &history);
            continue;
        }
        if harmony_message_id == Some(next) {
            if hstate == HarmonyState::ReadingChannel {
                // Channel name was accumulated but we don't need to parse it —
                // we just check via the channel_name buffer below
            }
            decode_buffer.clear();
            next = pick(&logits, sampling, &history);
            continue;
        }
        if hstate == HarmonyState::ReadingChannel {
            // Reading channel name tokens (e.g. "final", "analysis")
            let tok_text = tokenizer.decode(&[next]);
            if tok_text.contains("final") {
                hstate = HarmonyState::InFinal;
                in_thinking = false;
            } else {
                hstate = HarmonyState::InAnalysis;
                // Open a Qwen3-style thinking block for the analysis channel.
                if show_thinking {
                    print_stream_text("<think>\n", true, use_color);
                    in_thinking = true;
                }
            }
            next = pick(&logits, sampling, &history);
            continue;
        }
        if harmony_special.contains(&next) {
            next = pick(&logits, sampling, &history);
            continue;
        }
        if hstate == HarmonyState::InAnalysis {
            // Analysis channel = the model's reasoning. With --think, show it as a
            // thinking block (gray if color); otherwise suppress it (answer only).
            if show_thinking {
                print_generated_token(
                    tokenizer,
                    next,
                    &mut decode_buffer,
                    &mut in_thinking,
                    use_color,
                );
                io::stdout().flush().unwrap();
            }
            next = pick(&logits, sampling, &history);
            continue;
        }
        if is_moe && hstate != HarmonyState::InFinal {
            // Between harmony messages (after a channel's <|end|>, before the
            // next <|channel|>): the model emits a role header like "assistant".
            // That's structural, not user-visible content — suppress it. Only
            // for moe/harmony; non-moe (Qwen3) stays in Normal and prints here.
            next = pick(&logits, sampling, &history);
            continue;
        }

        answer_ids.push(next);
        print_generated_token(
            tokenizer,
            next,
            &mut decode_buffer,
            &mut in_thinking,
            use_color,
        );
        io::stdout().flush().unwrap();
        next = pick(&logits, sampling, &history);
    }

    print_stream_text(
        &tokenizer.flush_decode_stream(&mut decode_buffer),
        in_thinking,
        use_color,
    );
    if show_thinking && in_thinking {
        print_stream_text("\n</think>\n\n", true, use_color);
    }
    io::stdout().flush().unwrap();
    (Finish::Length, tokenizer.decode(&answer_ids))
}

fn append_after_stop(
    model: &ChatModel,
    cache: &mut PagedKVCache,
    tokenizer: &Tokenizer,
    max_seq_len: usize,
    _stop_token_id: u32,
    tp: &Option<TpHandle>,
) {
    append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n", tp);
}

fn append_text_to_cache(
    model: &ChatModel,
    cache: &mut PagedKVCache,
    tokenizer: &Tokenizer,
    max_seq_len: usize,
    text: &str,
    tp: &Option<TpHandle>,
) {
    let tokens = tokenizer.encode(text);
    if tokens.is_empty() || cache.seq_len(SLOT) + tokens.len() > max_seq_len {
        return;
    }
    if let Some(h) = tp {
        h.send(TpCommand::Prefill {
            tokens: tokens.clone(),
            slot: SLOT,
        });
    }
    let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
    if let Some(h) = tp {
        h.wait();
    }
}

fn print_generated_token(
    tokenizer: &Tokenizer,
    token_id: u32,
    decode_buffer: &mut Vec<u8>,
    in_thinking: &mut bool,
    use_color: bool,
) {
    if tokenizer.special_token_id("<think>") == Some(token_id) {
        print_stream_text(
            &tokenizer.flush_decode_stream(decode_buffer),
            *in_thinking,
            use_color,
        );
        *in_thinking = true;
        print_stream_text("<think>", true, use_color);
        return;
    }

    if tokenizer.special_token_id("</think>") == Some(token_id) {
        print_stream_text(
            &tokenizer.flush_decode_stream(decode_buffer),
            *in_thinking,
            use_color,
        );
        print_stream_text("</think>", true, use_color);
        *in_thinking = false;
        return;
    }

    let text = tokenizer.decode_token_stream(token_id, decode_buffer);
    print_stream_text(&text, *in_thinking, use_color);
}

fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
    if text.is_empty() {
        return;
    }
    if in_thinking && use_color {
        print!("\x1b[90m{text}\x1b[0m");
    } else {
        print!("{text}");
    }
}