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

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

    let model_dir = PathBuf::from(&args[1]);
    let max_tokens: usize = args
        .iter()
        .position(|a| a == "--max-tokens")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(100);

    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}).\n");

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

        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 next = sample_greedy_last(&logits);

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

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

                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,
                );
                next = sample_greedy_last(&logits);
            }
            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 next = match &model {
                Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
                Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
                Model::GptOss(_) => unreachable!(),
            };

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

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

                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!(),
                };
                next = match &model {
                    Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
                    Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
                    Model::GptOss(_) => unreachable!(),
                };
            }
            println!();
        }
    }
}

fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
    use half::bf16;
    assert_eq!(logits.ndim(), 2);
    let logits_cpu = logits.to_device(Device::Cpu);
    let vocab_size = logits.shape()[1];
    let seq_len = logits.shape()[0];
    let data = logits_cpu.as_slice::<bf16>();
    let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
    last.iter().enumerate()
        .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
        .map(|(i, _)| i as u32).unwrap()
}