mod inference;
mod prompt;
mod server;
mod stdio;

use std::io::Read;
use std::time::{Duration, Instant};

use clap::{Parser, Subcommand};
use inference::{InferenceConfig, InferenceEngine};

/// Errors for the sr-voice CLI.
#[derive(thiserror::Error, Debug)]
pub enum VoiceError {
    #[error("model load failed: {0}")]
    ModelLoadFailed(String),
    #[error("inference failed: {0}")]
    InferenceFailed(String),
    #[error("invalid input: {0}")]
    InvalidInput(String),
}

/// sr-voice — LLM inference service for The Settled Reach
#[derive(Parser)]
#[command(name = "sr-voice", version, about)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Start the inference server (loads model, listens for requests)
    Serve {
        /// Path to GGUF model file
        #[arg(long)]
        model: String,
        /// Listen port (ignored when --stdio is set)
        #[arg(long, default_value = "8321")]
        port: u16,
        /// CPU threads for inference
        #[arg(long)]
        threads: Option<u32>,
        /// Context window size in tokens
        #[arg(long, default_value = "512")]
        ctx_size: u32,
        /// Run in stdio mode: read JSONL from stdin, write JSONL to stdout.
        /// No network port is opened. Used by the game server's worker pool
        /// for Gemma 2 T&C compliance (no exposed inference endpoint).
        #[arg(long)]
        stdio: bool,
    },
    /// Generate text from a single prompt (requires running server)
    Generate {
        /// Server port
        #[arg(long, default_value = "8321")]
        port: u16,
        /// RNG seed
        #[arg(long)]
        seed: Option<u32>,
        /// Prompt file (reads from stdin if omitted)
        prompt_file: Option<String>,
    },
    /// Process a JSONL batch of prompts (requires running server)
    Batch {
        /// Server port
        #[arg(long, default_value = "8321")]
        port: u16,
        /// Input JSONL file
        #[arg(long)]
        input: String,
    },
    /// Run 5 inferences and report average tokens/sec (requires running server)
    Benchmark {
        /// Server port
        #[arg(long, default_value = "8321")]
        port: u16,
    },
}

fn default_threads() -> u32 {
    let cores = std::thread::available_parallelism()
        .map(|n| n.get() as u32)
        .unwrap_or(4);
    cores.saturating_sub(1).max(1)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();

    match cli.command {
        Command::Serve { model, port, threads, ctx_size, stdio } => {
            let threads = threads.unwrap_or_else(default_threads);
            let config = InferenceConfig {
                model_path: model.clone(),
                threads,
                ctx_size,
                seed: None,
            };

            eprintln!("Loading model: {}", config.model_path);
            let engine = InferenceEngine::load(&config)?;
            eprintln!("Model loaded ({} threads, {} ctx)", threads, ctx_size);

            if stdio {
                eprintln!("sr-voice stdio mode — reading JSONL from stdin");
                stdio::run_stdio(engine)?;
            } else {
                let model_name = std::path::Path::new(&model)
                    .file_name()
                    .map(|f| f.to_string_lossy().to_string())
                    .unwrap_or(model);

                server::run_server(engine, port, &model_name)?;
            }
        }
        Command::Generate { port, seed, prompt_file } => {
            let prompt = read_prompt(prompt_file)?;
            let req = serde_json::json!({ "prompt": prompt, "seed": seed });
            let body = post_with_status(port, "/generate", &req.to_string())?;
            let result: serde_json::Value = serde_json::from_str(&body)?;

            if let Some(err) = result.get("error") {
                return Err(format!("Server error: {}", err).into());
            }
            println!("{}", result["text"].as_str().unwrap_or(""));
            eprintln!(
                "[{} tokens in {}ms — {:.1} t/s, prefill {}ms]",
                result["tokens_generated"],
                result["generation_time_ms"],
                result["tokens_per_sec"].as_f64().unwrap_or(0.0),
                result["prefill_time_ms"],
            );
        }
        Command::Batch { port, input } => {
            let file = std::fs::File::open(&input)?;
            let reader = std::io::BufReader::new(file);
            let payloads = prompt::parse_jsonl(reader)?;

            let body = post_with_status(port, "/batch", &serde_json::to_string(&payloads)?)?;

            for line in body.lines() {
                if line.is_empty() { continue; }
                let result: serde_json::Value = serde_json::from_str(line)?;
                let id = result["id"].as_str().unwrap_or("?");
                if let Some(err) = result.get("error") {
                    eprintln!("--- {} --- ERROR: {}", id, err);
                } else {
                    println!("--- {} ---", id);
                    println!("{}", result["text"].as_str().unwrap_or(""));
                    eprintln!(
                        "[{} tokens in {}ms — {:.1} t/s]",
                        result["tokens_generated"],
                        result["generation_time_ms"],
                        result["tokens_per_sec"].as_f64().unwrap_or(0.0),
                    );
                }
            }
        }
        Command::Benchmark { port } => {
            let prompt = "Rephrase in terse dialect: The worker tends the crops in the field.";
            let runs = 5;
            eprintln!("Benchmark: {} runs", runs);

            let mut total_tps = 0.0;
            let mut total_prefill = 0u64;
            let mut total_gen = 0u64;

            for i in 0..runs {
                let req = serde_json::json!({ "prompt": prompt });
                let body = post_with_status(port, "/generate", &req.to_string())?;
                let result: serde_json::Value = serde_json::from_str(&body)?;

                let tps = result["tokens_per_sec"].as_f64().unwrap_or(0.0);
                let prefill = result["prefill_time_ms"].as_u64().unwrap_or(0);
                let gen = result["generation_time_ms"].as_u64().unwrap_or(0);
                let tokens = result["tokens_generated"].as_u64().unwrap_or(0);

                eprintln!("  run {}: {} tokens, {:.1} t/s, prefill {}ms", i + 1, tokens, tps, prefill);
                total_tps += tps;
                total_prefill += prefill;
                total_gen += gen;
            }

            eprintln!("\n=== Benchmark Results ===");
            eprintln!("  Avg tokens/sec:  {:.1}", total_tps / runs as f64);
            eprintln!("  Avg prefill:     {}ms", total_prefill / runs);
            eprintln!("  Avg generation:  {}ms", total_gen / runs);
        }
    }

    Ok(())
}

fn read_prompt(prompt_file: Option<String>) -> Result<String, Box<dyn std::error::Error>> {
    let raw = match prompt_file {
        Some(path) => std::fs::read_to_string(&path)?,
        None => {
            let mut buf = String::new();
            std::io::stdin().read_to_string(&mut buf)?;
            buf
        }
    };
    let trimmed = raw.trim().to_string();
    if trimmed.is_empty() {
        return Err("No prompt provided".into());
    }
    Ok(trimmed)
}

/// POST to the server. Prints "Server is processing..." if response takes > 500ms.
fn post_with_status(port: u16, path: &str, body: &str) -> Result<String, Box<dyn std::error::Error>> {
    let base = format!("http://127.0.0.1:{}", port);
    let agent = ureq::Agent::config_builder()
        .timeout_global(Some(Duration::from_secs(600)))
        .timeout_connect(Some(Duration::from_secs(2)))
        .build()
        .new_agent();

    // Health check — clear error if server isn't running
    if agent.get(&format!("{}/health", base)).call().is_err() {
        return Err(format!(
            "No sr-voice server on port {}. Start one with: sr-voice serve --model <path>",
            port
        ).into());
    }

    let url = format!("{}{}", base, path);
    let start = Instant::now();
    let printed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let flag = printed.clone();

    let handle = std::thread::spawn(move || {
        std::thread::sleep(Duration::from_millis(500));
        if !flag.load(std::sync::atomic::Ordering::Relaxed) {
            eprint!("Server is processing...");
            flag.store(true, std::sync::atomic::Ordering::Relaxed);
        }
    });

    let result = agent.post(&url)
        .header("Content-Type", "application/json")
        .send(body);

    let was_printed = printed.load(std::sync::atomic::Ordering::Relaxed);
    printed.store(true, std::sync::atomic::Ordering::Relaxed);
    let _ = handle.join();
    if was_printed {
        eprintln!(" done ({:.1}s)", start.elapsed().as_secs_f64());
    }

    match result {
        Ok(response) => {
            let text = response.into_body().read_to_string()?;
            Ok(text)
        }
        Err(e) => Err(format!("Request failed: {}", e).into()),
    }
}