use std::io::{self, Write};
use std::path::PathBuf;
use xserv_model::{
    BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, SamplingParams, loader, sample,
    sample_greedy_penalized,
};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;

fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
    args.iter()
        .position(|a| a == name)
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(default)
}

fn pick_next(
    logits: &xserv_tensor::Tensor,
    sampling: &SamplingParams,
    history: &[u32],
    rep_penalty: f32,
) -> u32 {
    if rep_penalty > 1.0 && sampling.temperature == 0.0 {
        sample_greedy_penalized(logits, history, rep_penalty)
    } else {
        sample(logits, sampling)
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!(
            "Usage: xserv-cli <model-dir> [--max-tokens N] [--temperature F] [--top-k N] [--top-p F] [--rep-penalty F] [--rep-window N]"
        );
        std::process::exit(1);
    }

    let model_dir = PathBuf::from(&args[1]);
    let max_tokens = flag(&args, "--max-tokens", 100usize);
    let sampling = SamplingParams {
        temperature: flag(&args, "--temperature", 0.0f32),
        top_k: flag(&args, "--top-k", 0usize),
        top_p: flag(&args, "--top-p", 1.0f32),
    };
    let rep_penalty = flag(&args, "--rep-penalty", 1.0f32);
    let rep_window = flag(&args, "--rep-window", 512usize);

    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(&model_dir.join("config.json"));
    let model_type = config.model_type.as_deref().unwrap_or("unknown");
    eprintln!(
        "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
        config.num_layers(),
        config.hidden(),
        config.num_heads(),
        config.num_kv_heads(),
        config.vocab_size
    );

    eprintln!("Loading weights...");
    let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
    eprintln!("Loaded {} tensors", weights.len());

    let is_qwen3 = model_type.contains("qwen");
    let is_gpt_oss = model_type.contains("gpt_oss");
    let dtype = if is_qwen3 || is_gpt_oss {
        DType::BF16
    } else {
        DType::F32
    };

    // Build model
    enum Model {
        GPT2(xserv_model::GPT2),
        Qwen3(xserv_model::Qwen3),
        GptOss(xserv_model::GptOss),
    }
    let model = if is_gpt_oss {
        Model::GptOss(xserv_model::GptOss::from_weights(config.clone(), weights))
    } else if is_qwen3 {
        Model::Qwen3(xserv_model::Qwen3::from_weights(config.clone(), weights))
    } else {
        Model::GPT2(xserv_model::GPT2::from_weights(config.clone(), weights))
    };

    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
    eprintln!(
        "Ready (KV cache, dtype={dtype}, temperature={}, top_k={}, top_p={}, rep_penalty={}, rep_window={}).\n",
        sampling.temperature, sampling.top_k, sampling.top_p, rep_penalty, rep_window
    );

    loop {
        print!("xserv> ");
        io::stdout().flush().unwrap();
        let mut input = String::new();
        if io::stdin().read_line(&mut input).unwrap() == 0 {
            break;
        }
        let raw_input = input.trim();
        if raw_input.is_empty() {
            continue;
        }
        if raw_input == "quit" || raw_input == "exit" {
            break;
        }
        let input = raw_input.replace("\\n", "\n");

        let token_ids = tokenizer.encode(&input);

        if is_gpt_oss {
            // GptOss uses paged KV cache
            let max_seq = 2048;
            let max_blocks_per_seq = (max_seq + BLOCK_SIZE - 1) / BLOCK_SIZE;
            let total_blocks = max_blocks_per_seq + 64;
            let mut paged_cache = PagedKVCache::new(
                &config,
                total_blocks,
                0,
                4,
                max_blocks_per_seq,
                DType::BF16,
                0,
            );
            let slot = 0;
            paged_cache.register_sequence(slot).expect("register slot");

            let model = match &model {
                Model::GptOss(m) => m,
                _ => unreachable!(),
            };
            let logits = model.forward_prefill_paged(&token_ids, slot, &mut paged_cache);
            let mut history = token_ids.clone();
            let start = history.len().saturating_sub(rep_window);
            let mut next = pick_next(&logits, &sampling, &history[start..], rep_penalty);

            print!("{input}");
            io::stdout().flush().unwrap();

            for _ in 0..max_tokens {
                let text = tokenizer.decode(&[next]);
                print!("{text}");
                io::stdout().flush().unwrap();
                history.push(next);

                if tokenizer.eos_token_id() == Some(next) {
                    break;
                }

                let pos = paged_cache.seq_len(slot);
                let logits = model.forward_decode_paged(&[next], &[pos], &[slot], &mut paged_cache);
                let start = history.len().saturating_sub(rep_window);
                next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
            }
            println!();
            paged_cache.free_sequence(slot);
        } else {
            let kv_heads = if is_qwen3 {
                config.num_kv_heads()
            } else {
                config.num_heads()
            };
            let mut cache = KVCache::new(
                config.num_layers(),
                kv_heads,
                config.head_dim(),
                dtype,
                Device::Cuda(0),
            );

            let logits = match &model {
                Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
                Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
                Model::GptOss(_) => unreachable!(),
            };
            let mut history = token_ids.clone();
            let start = history.len().saturating_sub(rep_window);
            let mut next = pick_next(&logits, &sampling, &history[start..], rep_penalty);

            print!("{input}");
            io::stdout().flush().unwrap();

            for _ in 0..max_tokens {
                let text = tokenizer.decode(&[next]);
                print!("{text}");
                io::stdout().flush().unwrap();
                history.push(next);

                if tokenizer.eos_token_id() == Some(next) {
                    break;
                }

                let logits = match &model {
                    Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
                    Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
                    Model::GptOss(_) => unreachable!(),
                };
                let start = history.len().saturating_sub(rep_window);
                next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
            }
            println!();
        }
    }
}