//! Stitched ad presentation (extensions/stitched-presentation.md): the corpus, the verifier
//! and the pod-input signer, as a tool and as a service (§8).
//!
//!   wai_stitched_presentation vectors gen    <corpus-dir>
//!   wai_stitched_presentation vectors verify <corpus-dir>
//!       write live-conformance/stitched/ from fixed seeds; or check that every file
//!       regenerates byte for byte, that no other file but README.md is there, and that
//!       every vector gets the facts it states
//!   wai_stitched_presentation verify <bundle.json>
//!       verify a bundle (records, claims, manifests, responses, receipts, creatives,
//!       issuers, signers; README of live-conformance/stitched/) and print what §7.5 states
//!   wai_stitched_presentation sign-pod <draft.json> --key <seed-file>
//!       write and sign a pod-input record (§4) from a draft and print it in hexadecimal;
//!       the seed file holds the signer's 32-byte Ed25519 seed in hexadecimal
//!   wai_stitched_presentation serve [--listen 127.0.0.1:9690] [--max-body BYTES]
//!                                   [--max-connections N] [--read-timeout-ms MS]
//!       answer `POST /v1/verify` (a bundle in, the statement out) and `GET /v1/health`
//!       over HTTP/1.1 on a loopback address, one request a connection
//!
//! The corpus's second implementation is `live-conformance/stitched_presentation_verify.py`,
//! written from the extension's text and sharing no code with this binary.

use std::collections::BTreeMap;
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use ed25519_dalek::SigningKey;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use wai::live_ads::presentation::{Creative, CreativeObject, Issuer, Requirement};
use wai::stitched_presentation::{self as sp, corpus, Frame, Marker, PodAd, PodInput};
use wai::wai_video_ssai::terms::Consent;

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let code = match a.get(1).map(String::as_str) {
        Some("vectors") => match (a.get(2).map(String::as_str), a.get(3)) {
            (Some("gen"), Some(d)) => vectors_gen(Path::new(d)),
            (Some("verify"), Some(d)) => vectors_verify(Path::new(d)),
            _ => usage(),
        },
        Some("verify") => match a.get(2) {
            Some(p) => match std::fs::read(p).map_err(|e| e.to_string()).and_then(|b| statement(&b)) {
                Ok(v) => {
                    println!("{}", serde_json::to_string_pretty(&v).expect("JSON"));
                    0
                }
                Err(e) => {
                    eprintln!("verify: {e}");
                    1
                }
            },
            None => usage(),
        },
        Some("sign-pod") => match (a.get(2), opt(&a, "--key")) {
            (Some(draft), Some(key)) => match sign_pod(Path::new(draft), Path::new(key)) {
                Ok(hex_rec) => {
                    println!("{hex_rec}");
                    0
                }
                Err(e) => {
                    eprintln!("sign-pod: {e}");
                    1
                }
            },
            _ => usage(),
        },
        Some("serve") => serve(&a[2..]),
        _ => usage(),
    };
    std::process::exit(code);
}

fn usage() -> i32 {
    eprintln!("usage: wai_stitched_presentation vectors gen|verify <dir> | verify <bundle.json> | sign-pod <draft.json> --key <seed-file> | serve [--listen ADDR] [--max-body N] [--max-connections N] [--read-timeout-ms MS]");
    2
}

fn opt<'a>(a: &'a [String], flag: &str) -> Option<&'a str> {
    a.iter().position(|x| x == flag).and_then(|i| a.get(i + 1)).map(String::as_str)
}

// ---- the corpus ----

fn vectors_gen(d: &Path) -> i32 {
    if let Err(e) = std::fs::create_dir_all(d) {
        eprintln!("create {}: {e}", d.display());
        return 1;
    }
    for (name, body) in corpus::generate() {
        if let Err(e) = std::fs::write(d.join(&name), body) {
            eprintln!("write {name}: {e}");
            return 1;
        }
    }
    println!("wrote {}", d.display());
    0
}

fn vectors_verify(d: &Path) -> i32 {
    let mut files = BTreeMap::new();
    let Ok(entries) = std::fs::read_dir(d) else {
        eprintln!("read {}", d.display());
        return 1;
    };
    for e in entries.flatten() {
        let name = e.file_name().to_string_lossy().into_owned();
        if name == "README.md" {
            continue;
        }
        match std::fs::read_to_string(e.path()) {
            Ok(s) => {
                files.insert(name, s);
            }
            Err(err) => {
                eprintln!("read {name}: {err}");
                return 1;
            }
        }
    }
    let generated: BTreeMap<String, String> = corpus::generate().into_iter().collect();
    let same = files == generated;
    println!("{} every file regenerates byte for byte from its seeds, and no other file is there", if same { "PASS" } else { "FAIL" });
    let fails = corpus::verify(&files);
    for f in &fails {
        println!("FAIL {f}");
    }
    println!("{} every vector gets the facts the corpus states", if fails.is_empty() { "PASS" } else { "FAIL" });
    (!same || !fails.is_empty()) as i32
}

// ---- a bundle ----

fn unhex(s: &str) -> Result<Vec<u8>, String> {
    if !s.len().is_multiple_of(2) || !s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) {
        return Err(format!("not lowercase hexadecimal: {s:.40}"));
    }
    Ok((0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).expect("hex")).collect())
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

fn hex_list(v: &Value, k: &str) -> Result<Vec<Vec<u8>>, String> {
    v.get(k).and_then(Value::as_array).ok_or(format!("{k}: an array"))?.iter().map(|x| unhex(x.as_str().ok_or(format!("{k}: hexadecimal strings"))?)).collect()
}

fn text_list(v: &Value, k: &str) -> Result<Vec<Vec<u8>>, String> {
    match v.get(k) {
        None => Ok(Vec::new()),
        Some(a) => a.as_array().ok_or(format!("{k}: an array"))?.iter().map(|x| x.as_str().map(|s| s.as_bytes().to_vec()).ok_or(format!("{k}: strings"))).collect(),
    }
}

/// Verify a bundle and state the result (§7.5).
fn statement(b: &[u8]) -> Result<Value, String> {
    let v: Value = serde_json::from_slice(b).map_err(|e| format!("a JSON object: {e}"))?;
    let records = hex_list(&v, "records")?;
    let claims = hex_list(&v, "claims")?;
    if claims.len() > sp::MAX_CLAIMS {
        return Err(format!("at most {} claims a verification", sp::MAX_CLAIMS));
    }
    let manifests = text_list(&v, "manifests")?;
    let responses = text_list(&v, "responses")?;
    let receipts: Vec<Value> = v.get("receipts").and_then(Value::as_array).cloned().unwrap_or_default();
    let creatives: Vec<Creative> = v
        .get("creatives")
        .and_then(Value::as_array)
        .ok_or("creatives: an array")?
        .iter()
        .map(|c| {
            let objs = c.get("objects").and_then(Value::as_array).ok_or("a creative's objects")?;
            let objects = objs
                .iter()
                .map(|o| {
                    let h: [u8; 32] = unhex(o.get("content_hash").and_then(Value::as_str).ok_or("content_hash")?)?.try_into().map_err(|_| "a 32-byte content hash")?;
                    let d = u32::try_from(o.get("duration_ms").and_then(Value::as_u64).ok_or("duration_ms")?).map_err(|_| "duration_ms")?;
                    Ok(CreativeObject { content_hash: h, duration_ms: d })
                })
                .collect::<Result<Vec<_>, String>>()?;
            Ok(Creative { objects })
        })
        .collect::<Result<_, String>>()?;
    let issuers: Vec<Issuer> = match v.get("issuers").and_then(Value::as_array) {
        None => Vec::new(),
        Some(a) => a
            .iter()
            .map(|i| {
                Ok(Issuer {
                    name: i.get("name").and_then(Value::as_str).ok_or("an issuer's name")?.as_bytes().to_vec(),
                    token_type: u16::try_from(i.get("token_type").and_then(Value::as_u64).ok_or("token_type")?).map_err(|_| "token_type")?,
                    keys: hex_list(i, "keys")?,
                })
            })
            .collect::<Result<_, String>>()?,
    };
    let signers: Vec<[u8; 32]> = hex_list(&v, "signers")?.into_iter().map(|k| k.try_into().map_err(|_| "a 32-byte signer key".to_string())).collect::<Result<_, _>>()?;
    let got = sp::verify(&sp::Input { records: &records, claims: &claims, manifests: &manifests, responses: &responses, receipts: &receipts, creatives: &creatives, issuers: &issuers, signers: &signers });
    let mut o = corpus::facts_json(&got);
    o.insert("capability".into(), json!(sp::CAPABILITY));
    Ok(Value::Object(o))
}

// ---- the pod-input signer ----

fn field<'a>(v: &'a Value, k: &str) -> Result<&'a Value, String> {
    v.get(k).ok_or(format!("{k}: missing"))
}

fn num(v: &Value, k: &str) -> Result<u64, String> {
    field(v, k)?.as_u64().ok_or(format!("{k}: an integer"))
}

fn text(v: &Value, k: &str) -> Result<Vec<u8>, String> {
    Ok(field(v, k)?.as_str().ok_or(format!("{k}: a string"))?.as_bytes().to_vec())
}

fn marker_of(carriage: u8, m: &Value) -> Result<Marker, String> {
    let s = |k: &str| -> Result<Vec<u8>, String> { Ok(m.get(k).and_then(Value::as_str).unwrap_or("").as_bytes().to_vec()) };
    Ok(match carriage {
        sp::carriage::HLS_DATERANGE => Marker::DateRange { id: text(m, "id")?, start: text(m, "start")? },
        sp::carriage::HLS_INTERSTITIAL => Marker::Interstitial { id: text(m, "id")?, start: text(m, "start")?, index: u16::try_from(num(m, "index")?).map_err(|_| "index")? },
        sp::carriage::DASH_EVENT => Marker::Event { period: s("period")?, scheme: text(m, "scheme")?, value: s("value")?, id: text(m, "id")? },
        _ => return Err("carriage: 1, 2 or 3".into()),
    })
}

/// A draft (README of live-conformance/stitched/): the record's fields, a marker's fields in
/// place of its digest, and the response by its file or its digest.
fn sign_pod(draft: &Path, key: &Path) -> Result<String, String> {
    let d: Value = serde_json::from_slice(&std::fs::read(draft).map_err(|e| e.to_string())?).map_err(|e| e.to_string())?;
    let seed = String::from_utf8(std::fs::read(key).map_err(|e| format!("the key file: {e}"))?).map_err(|_| "the key file: hexadecimal")?;
    let seed: [u8; 32] = unhex(seed.trim())?.try_into().map_err(|_| "the key file: a 32-byte seed")?;
    let carriage = u8::try_from(num(&d, "carriage")?).map_err(|_| "carriage")?;
    let response: [u8; 32] = match (d.get("response_file").and_then(Value::as_str), d.get("response").and_then(Value::as_str)) {
        (Some(f), _) => {
            let p = draft.parent().unwrap_or(Path::new(".")).join(f);
            Sha256::digest(std::fs::read(&p).map_err(|e| format!("{}: {e}", p.display()))?).into()
        }
        (None, Some(h)) => unhex(h)?.try_into().map_err(|_| "response: 32 bytes")?,
        _ => return Err("response_file or response".into()),
    };
    let session: [u8; 16] = match d.get("session").and_then(Value::as_str) {
        None => [0; 16],
        Some(h) => unhex(h)?.try_into().map_err(|_| "session: 16 bytes")?,
    };
    let f = d.get("frame").cloned().unwrap_or(Value::Null);
    let fnum = |k: &str| f.get(k).and_then(Value::as_u64).unwrap_or(0);
    let frame = Frame {
        kind: fnum("kind") as u8,
        width: fnum("width") as u16,
        height: fnum("height") as u16,
        corner: fnum("corner") as u8,
        window_w: fnum("window_w") as u16,
        window_h: fnum("window_h") as u16,
        gap: fnum("gap") as u16,
        aspect: fnum("aspect") as u8,
    };
    let r = field(&d, "requirement")?;
    let requirement = Requirement {
        require: num(r, "require")? as u8,
        rule: num(r, "rule")? as u8,
        token_type: num(r, "token_type")? as u16,
        issuer: r.get("issuer").and_then(Value::as_str).unwrap_or("").as_bytes().to_vec(),
    };
    let currency: [u8; 3] = text(&d, "currency")?.try_into().map_err(|_| "currency: three letters")?;
    // A draft's `consent` (§4.7) makes a record of version 2: its GPP string, its applicable
    // sections and its child-directed flag, as live-ads §8.6's JSON carriage writes them.
    let consent = match d.get("consent") {
        None => None,
        Some(c) => Some(Consent::from_json(c).ok_or("consent: {\"gpp\", \"gpp_sid\", \"child_directed\"}")?),
    };
    let ads = field(&d, "ads")?
        .as_array()
        .ok_or("ads: an array")?
        .iter()
        .enumerate()
        .map(|(k, a)| {
            let evidence: [u8; 32] = match a.get("evidence").and_then(Value::as_str) {
                None => [0; 32],
                Some(h) => unhex(h)?.try_into().map_err(|_| "evidence: 32 bytes")?,
            };
            let creative_hash: [u8; 32] = unhex(a.get("creative_hash").and_then(Value::as_str).ok_or("creative_hash")?)?.try_into().map_err(|_| "creative_hash: 32 bytes")?;
            Ok(PodAd {
                slot: k as u8,
                marker: marker_of(carriage, field(a, "marker")?)?.digest(),
                ad_id: text(a, "ad_id")?,
                serving_id: text(a, "serving_id")?,
                creative_id: a.get("creative_id").and_then(Value::as_str).unwrap_or("").as_bytes().to_vec(),
                registry: a.get("registry").and_then(Value::as_str).unwrap_or("").as_bytes().to_vec(),
                creative_hash,
                duration_ms: u32::try_from(num(a, "duration_ms")?).map_err(|_| "duration_ms")?,
                price_micros: num(a, "price_micros")?,
                price_source: num(a, "price_source")? as u8,
                evidence,
            })
        })
        .collect::<Result<Vec<_>, String>>()?;
    let rec = PodInput {
        carriage,
        session,
        avail_id: text(&d, "avail_id")?,
        avail_ms: u32::try_from(num(&d, "avail_ms")?).map_err(|_| "avail_ms")?,
        response_format: num(&d, "response_format")? as u8,
        response,
        id_kind: num(&d, "id_kind")? as u8,
        id_attr: d.get("id_attr").and_then(Value::as_str).unwrap_or("").as_bytes().to_vec(),
        currency,
        exponent: num(&d, "exponent")? as u8,
        frame,
        requirement,
        ads,
        consent,
        signer: [0; 32],
        sig: [0; 64],
    };
    if !rec.rules_hold() {
        return Err("the draft breaks a rule of §4.4 (pod_input_malformed)".into());
    }
    Ok(hex(&rec.sign(&SigningKey::from_bytes(&seed)).to_bytes()))
}

// ---- the service ----

struct Config {
    max_body: usize,
    read_timeout: Duration,
}

const MAX_HEAD: usize = 16 * 1024;

fn respond(mut s: &TcpStream, status: u16, body: &Value) {
    let reason = match status {
        200 => "OK",
        400 => "Bad Request",
        404 => "Not Found",
        405 => "Method Not Allowed",
        408 => "Request Timeout",
        413 => "Content Too Large",
        431 => "Request Header Fields Too Large",
        503 => "Service Unavailable",
        _ => "Error",
    };
    let b = serde_json::to_vec(body).expect("JSON");
    let head = format!("HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nCache-Control: no-store\r\nConnection: close\r\n\r\n", b.len());
    let _ = s.write_all(head.as_bytes()).and_then(|_| s.write_all(&b));
}

fn handle(cfg: &Config, s: TcpStream) {
    let deadline = Instant::now() + cfg.read_timeout;
    let mut buf = Vec::new();
    let mut chunk = [0u8; 8192];
    let head_end = loop {
        if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            break i + 4;
        }
        if buf.len() > MAX_HEAD {
            return respond(&s, 431, &json!({"error": "request head too large"}));
        }
        let left = deadline.saturating_duration_since(Instant::now());
        if left.is_zero() || s.set_read_timeout(Some(left)).is_err() {
            return respond(&s, 408, &json!({"error": "request head not received in time"}));
        }
        match (&s).read(&mut chunk) {
            Ok(0) | Err(_) => return respond(&s, 408, &json!({"error": "request head not received in time"})),
            Ok(n) => buf.extend_from_slice(&chunk[..n]),
        }
    };
    let head = String::from_utf8_lossy(&buf[..head_end]).into_owned();
    let mut lines = head.split("\r\n");
    let mut first = lines.next().unwrap_or("").split(' ');
    let (method, path) = (first.next().unwrap_or(""), first.next().unwrap_or(""));
    let length = lines.filter_map(|l| l.split_once(':')).find(|(k, _)| k.eq_ignore_ascii_case("content-length")).and_then(|(_, v)| v.trim().parse::<usize>().ok());
    match (method, path) {
        ("GET", "/v1/health") => respond(&s, 200, &json!({"ok": true, "capability": sp::CAPABILITY})),
        ("POST", "/v1/verify") => {
            let Some(n) = length else { return respond(&s, 400, &json!({"error": "Content-Length required"})) };
            if n > cfg.max_body {
                return respond(&s, 413, &json!({"error": format!("a bundle of at most {} bytes", cfg.max_body)}));
            }
            let mut body = buf[head_end..].to_vec();
            while body.len() < n {
                let left = deadline.saturating_duration_since(Instant::now());
                if left.is_zero() || s.set_read_timeout(Some(left)).is_err() {
                    return respond(&s, 408, &json!({"error": "request body not received in time"}));
                }
                match (&s).read(&mut chunk) {
                    Ok(0) | Err(_) => return respond(&s, 408, &json!({"error": "request body not received in time"})),
                    Ok(k) => body.extend_from_slice(&chunk[..k]),
                }
            }
            body.truncate(n);
            match statement(&body) {
                Ok(v) => respond(&s, 200, &v),
                Err(e) => respond(&s, 400, &json!({"error": e})),
            }
        }
        (_, "/v1/health" | "/v1/verify") => respond(&s, 405, &json!({"error": "method not allowed"})),
        _ => respond(&s, 404, &json!({"error": "not found"})),
    }
}

fn serve(a: &[String]) -> i32 {
    let listen = opt(a, "--listen").unwrap_or("127.0.0.1:9690");
    let addr: SocketAddr = match listen.parse() {
        Ok(x) => x,
        Err(_) => {
            eprintln!("--listen: an address and port");
            return 2;
        }
    };
    if !addr.ip().is_loopback() {
        eprintln!("--listen: a loopback address; a reverse proxy publishes the service");
        return 2;
    }
    let num = |k: &str, d: usize| opt(a, k).map_or(Some(d), |v| v.parse().ok());
    let (Some(max_body), Some(max_conn), Some(read_ms)) = (num("--max-body", 8 << 20), num("--max-connections", 32), num("--read-timeout-ms", 10_000)) else {
        eprintln!("--max-body, --max-connections and --read-timeout-ms take integers");
        return 2;
    };
    let listener = match TcpListener::bind(addr) {
        Ok(l) => l,
        Err(e) => {
            eprintln!("bind {addr}: {e}");
            return 1;
        }
    };
    eprintln!("wai_stitched_presentation: serving on http://{addr}");
    let cfg = Arc::new(Config { max_body, read_timeout: Duration::from_millis(read_ms as u64) });
    let open = Arc::new(AtomicUsize::new(0));
    for stream in listener.incoming().flatten() {
        if open.load(Ordering::SeqCst) >= max_conn {
            respond(&stream, 503, &json!({"error": "busy"}));
            continue;
        }
        open.fetch_add(1, Ordering::SeqCst);
        let (cfg, open) = (Arc::clone(&cfg), Arc::clone(&open));
        std::thread::spawn(move || {
            handle(&cfg, stream);
            open.fetch_sub(1, Ordering::SeqCst);
        });
    }
    0
}
