//! Minimal dependency-free HTTP/1.1 API for token-ID-based Falcon-H1 inference.
//! For localhost/Colab tunnel demo only; deploy behind a reverse proxy for production.
use falcon_h1_rust::{load_pack,synthetic_weights,Config,Model};
use std::io::{Read,Write};
use std::net::{TcpListener,TcpStream};
use std::sync::Arc;
use std::time::Duration;

fn extract_array(body:&str,key:&str)->Result<Vec<usize>,String>{
    let needle=format!("\"{key}\"");
    let idx=body.find(&needle).ok_or_else(||format!("missing {key}"))?;
    let tail=&body[idx+needle.len()..];
    let colon=tail.find(':').ok_or("missing :")?;
    let tail=&tail[colon+1..];
    let start=tail.find('[').ok_or("missing [")?;
    let tail=&tail[start+1..];
    let end=tail.find(']').ok_or("missing ]")?;
    let data=tail[..end].trim();
    if data.is_empty(){return Err("empty input_ids".into())}
    let ids=data.split(',').map(|x|x.trim().parse::<usize>().map_err(|_|"invalid token ID".to_string())).collect::<Result<Vec<_>,_>>()?;
    if ids.len()>2048{return Err("sequence exceeds API limit (2048 tokens)".into())}
    Ok(ids)
}
fn extract_number(body:&str,key:&str,default:usize)->Result<usize,String>{
    let needle=format!("\"{key}\"");
    if let Some(idx)=body.find(&needle){
        let tail=&body[idx+needle.len()..];
        let colon=tail.find(':').ok_or("missing :")?;
        let num=tail[colon+1..].trim_start().chars().take_while(|c|c.is_ascii_digit()).collect::<String>();
        if num.is_empty(){return Err(format!("invalid {key}"))}
        num.parse::<usize>().map_err(|_|format!("invalid {key}"))
    } else {Ok(default)}
}
fn response(mut stream:TcpStream,code:&str,payload:&str){
    let msg=format!("HTTP/1.1 {code}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\nAccess-Control-Allow-Origin: *\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type\r\n\r\n{}",payload.len(),payload);
    let _=stream.write_all(msg.as_bytes());
}
fn json_err(e:&str)->String{format!("{{\"error\":\"{}\"}}",e.replace('\\',"\\\\").replace('"',"\\\"").replace('\n'," "))}
enum Backend {
    Cpu(Model),
    #[cfg(feature="candle")]
    Candle(falcon_h1_rust::candle_backend::CandleModel),
}
impl Backend {
    fn config(&self)->&Config { match self {
        Self::Cpu(m)=>&m.config,
        #[cfg(feature="candle")]
        Self::Candle(m)=>&m.config,
    }}
    fn label(&self)->&'static str {match self {
        Self::Cpu(_)=>"cpu-reference",
        #[cfg(feature="candle")]
        Self::Candle(m)=>m.device_name(),
    }}
    fn logits(&self,ids:&[usize])->Result<Vec<f32>,String>{match self {
        Self::Cpu(m)=>m.logits(ids),
        #[cfg(feature="candle")]
        Self::Candle(m)=>m.logits(ids),
    }}
    fn generate(&self,ids:&[usize],n:usize)->Result<Vec<usize>,String>{match self {
        Self::Cpu(m)=>m.generate(ids,n),
        #[cfg(feature="candle")]
        Self::Candle(m)=>m.generate(ids,n),
    }}
}
fn handle(mut stream:TcpStream,model:Arc<Backend>){
    let _=stream.set_read_timeout(Some(Duration::from_secs(15)));
    let _=stream.set_write_timeout(Some(Duration::from_secs(15)));
    let mut buf=vec![0u8;16384];let mut len=0usize;let mut expected=None;
    loop {
        if len==buf.len(){response(stream,"413 Payload Too Large",&json_err("request too large"));return}
        match stream.read(&mut buf[len..]) {
            Ok(0)=>break,
            Ok(n)=>len+=n,
            Err(_)=>return,
        }
        if let Some(hend)=buf[..len].windows(4).position(|w|w==b"\r\n\r\n") {
            if expected.is_none(){
                let headers=String::from_utf8_lossy(&buf[..hend]);
                let size=headers.lines().filter_map(|l|l.split_once(':')).find(|(k,_)|k.eq_ignore_ascii_case("content-length"))
                    .and_then(|(_,v)|v.trim().parse::<usize>().ok()).unwrap_or(0);
                expected=Some(hend+4+size);
            }
            if len>=expected.unwrap(){break}
        }
    }
    let raw=String::from_utf8_lossy(&buf[..len]);
    let first=raw.lines().next().unwrap_or("");
    let (method,path) = {let mut p=first.split_whitespace();(p.next().unwrap_or(""),p.next().unwrap_or(""))};
    if method=="OPTIONS" {response(stream,"204 No Content","");return}
    if method=="GET" && path=="/health" {response(stream,"200 OK",&format!("{{\"status\":\"ok\",\"backend\":\"{}\"}}",model.label()));return}
    if method=="GET" && path=="/v1/model" {
        let c=model.config();
        let msg=format!("{{\"architecture\":\"falcon_h1\",\"backend\":\"{}\",\"layers\":{},\"hidden\":{},\"attention_heads\":{},\"kv_heads\":{},\"mamba_heads\":{},\"vocab_size\":{}}}",model.label(),c.layers,c.hidden,c.heads,c.kv_heads,c.mamba_heads,c.vocab);
        response(stream,"200 OK",&msg);return
    }
    if method!="POST" || (path!="/v1/logits" && path!="/v1/generate") {
        response(stream,"404 Not Found",&json_err("unknown route"));return
    }
    let body=match raw.split_once("\r\n\r\n") {Some((_,v))=>v,None=>{response(stream,"400 Bad Request",&json_err("no body"));return}};
    let result = (||->Result<String,String>{
        let ids=extract_array(body,"input_ids")?;
        if path=="/v1/logits" {
            let y=model.logits(&ids)?;
            let text=y.iter().map(|x|format!("{x:.7}")).collect::<Vec<_>>().join(",");
            Ok(format!("{{\"logits\":[{text}],\"vocab_size\":{}}}",y.len()))
        } else {
            let count=extract_number(body,"max_new_tokens",1)?;
            let seq=model.generate(&ids,count)?;
            Ok(format!("{{\"output_ids\":[{}],\"generated_tokens\":{count}}}",seq.iter().map(|x|x.to_string()).collect::<Vec<_>>().join(",")))
        }
    })();
    match result{Ok(msg)=>response(stream,"200 OK",&msg),Err(e)=>response(stream,"400 Bad Request",&json_err(&e))}
}
fn main()->Result<(),Box<dyn std::error::Error>>{
    let args:Vec<String>=std::env::args().collect();
    let demo=args.iter().any(|x|x=="--demo");
    let bind=args.iter().position(|x|x=="--bind").and_then(|i|args.get(i+1)).map(|s|s.as_str()).unwrap_or("127.0.0.1:8080");
    let c=if demo{Config::demo()}else{Config::falcon_h1_05b()};
    let weights=if demo{synthetic_weights(&c)}else{
        let path=args.iter().position(|x|x=="--checkpoint").and_then(|i|args.get(i+1)).ok_or("use --checkpoint path/to/fh1.f32pack or --demo")?;
        load_pack(path)?
    };
    let device=args.iter().position(|x|x=="--device").and_then(|i|args.get(i+1)).map(|s|s.as_str()).unwrap_or("cpu");
    let model=match device {
        "cpu"=>Backend::Cpu(Model::new(c,weights)?),
        #[cfg(feature="candle")]
        "cpu-candle"=>Backend::Candle(falcon_h1_rust::candle_backend::CandleModel::new(c,weights,candle_core::Device::Cpu)?),
        #[cfg(feature="cuda")]
        "cuda"=>Backend::Candle(falcon_h1_rust::candle_backend::CandleModel::new(c,weights,candle_core::Device::new_cuda(0)?)?),
        _=>return Err(format!("unsupported --device {device}; build with --features cuda for CUDA").into()),
    };
    let model=Arc::new(model);
    let listener=TcpListener::bind(bind)?;
    eprintln!("Falcon-H1 Rust API listening on http://{bind} (demo={demo}, backend={})",model.label());
    for request in listener.incoming(){
        if let Ok(stream)=request {let shared=Arc::clone(&model);
            std::thread::spawn(move||handle(stream,shared));}
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test] fn parse_ids(){ assert_eq!(extract_array(r#"{"input_ids":[1,2,3]}"#,"input_ids").unwrap(),vec![1,2,3]); }
    #[test] fn reject_ids(){assert!(extract_array(r#"{"input_ids":[1,"x"]}"#,"input_ids").is_err());}
    #[test] fn parse_count(){ assert_eq!(extract_number(r#"{"max_new_tokens":3}"#,"max_new_tokens",1).unwrap(),3); }
}