use std::io::{self, Write};
use std::path::PathBuf;
use xserv_model::{loader, KVCache, ModelConfig};
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 dtype = if is_qwen3 { DType::BF16 } else { DType::F32 };

    // Build model
    enum Model {
        GPT2(xserv_model::GPT2),
        Qwen3(xserv_model::Qwen3),
    }
    let model = 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);
        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),
        );

        // Prefill + decode
        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),
        };
        let mut next = match &model {
            Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
            Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&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 logits = match &model {
                Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
                Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
            };
            next = match &model {
                Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
                Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
            };
        }
        println!();
    }
}