use anyhow::{Context, Result, bail, ensure};
use ort::{session::Session, value::Tensor};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::{
    io::{self, BufRead, Write},
    path::Path,
};
use tokenizers::{Encoding, Tokenizer};

#[derive(Deserialize)]
struct Config {
    labels: Vec<String>,
    max_length: usize,
    stride: usize,
    #[serde(default)]
    calibration: Option<Calibration>,
    #[serde(default)]
    calibration_file: Option<String>,
}
#[derive(Deserialize)]
struct Calibration {
    #[serde(alias = "xs")]
    x: Vec<f64>,
    #[serde(alias = "ys")]
    y: Vec<f64>,
    threshold: Option<f64>,
}
impl Calibration {
    fn apply(&self, v: f64) -> f64 {
        let i = self.x.partition_point(|x| *x < v);
        if i == 0 {
            self.y[0]
        } else if i == self.x.len() {
            *self.y.last().unwrap()
        } else {
            self.y[i - 1]
                + (v - self.x[i - 1]) / (self.x[i] - self.x[i - 1]) * (self.y[i] - self.y[i - 1])
        }
    }
}
#[derive(Clone, Debug, Serialize)]
struct Finding {
    label: String,
    start: usize,
    end: usize,
    raw_score: f64,
    confidence: Option<f64>,
    #[serde(skip)]
    boundary_fragment: bool,
}
fn encode_windows(
    tokenizer: &Tokenizer,
    text: &str,
    max_length: usize,
    stride: usize,
) -> Result<Vec<Encoding>> {
    let full = tokenizer
        .encode(text, false)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let capacity = max_length - 2;
    ensure!(stride < capacity, "invalid stride");
    let cls = tokenizer
        .token_to_id("[CLS]")
        .context("missing CLS token")?;
    let sep = tokenizer
        .token_to_id("[SEP]")
        .context("missing SEP token")?;
    let mut windows = Vec::new();
    let mut start = 0;
    loop {
        let end = (start + capacity).min(full.len());
        let n = end - start + 2;
        let mut ids = vec![cls];
        ids.extend_from_slice(&full.get_ids()[start..end]);
        ids.push(sep);
        let mut offsets = vec![(0, 0)];
        offsets.extend_from_slice(&full.get_offsets()[start..end]);
        offsets.push((0, 0));
        let mut tokens = vec!["[CLS]".to_string()];
        tokens.extend_from_slice(&full.get_tokens()[start..end]);
        tokens.push("[SEP]".to_string());
        let mut special = vec![0; n];
        special[0] = 1;
        special[n - 1] = 1;
        windows.push(Encoding::new(
            ids,
            vec![0; n],
            tokens,
            vec![None; n],
            offsets,
            special,
            vec![1; n],
            vec![],
            Default::default(),
        ));
        if end == full.len() {
            break;
        }
        start += capacity - stride;
    }
    Ok(windows)
}
fn decode(enc: &tokenizers::Encoding, logits: &[f32], labels: &[String]) -> Vec<Finding> {
    let k = 1 + 2 * labels.len();
    let mut out = Vec::new();
    let mut current: Option<Finding> = None;
    for (i, row) in logits.chunks_exact(k).enumerate() {
        let (start, end) = enc.get_offsets()[i];
        if enc.get_special_tokens_mask()[i] != 0 || enc.get_attention_mask()[i] == 0 || start == end
        {
            if let Some(f) = current.take() {
                out.push(f);
            }
            continue;
        }
        let mut id = 0;
        for j in 1..k {
            if row[j] > row[id] {
                id = j;
            }
        }
        if id == 0 {
            if let Some(f) = current.take() {
                out.push(f);
            }
            continue;
        }
        let label = &labels[(id - 1) / 2];
        let max = row[id] as f64;
        let prob = 1.0 / row.iter().map(|v| (*v as f64 - max).exp()).sum::<f64>();
        if id % 2 == 0 && current.as_ref().is_some_and(|f| f.label == *label) {
            let f = current.as_mut().unwrap();
            f.end = end;
            f.raw_score = f.raw_score.min(prob);
        } else {
            if let Some(f) = current.take() {
                out.push(f);
            }
            current = Some(Finding {
                label: label.clone(),
                start,
                end,
                raw_score: prob,
                confidence: None,
                boundary_fragment: false,
            });
        }
    }
    if let Some(f) = current {
        out.push(f);
    }
    out
}
fn mark_boundaries(
    enc: &tokenizers::Encoding,
    findings: &mut [Finding],
    has_previous: bool,
    has_next: bool,
) {
    let content: Vec<_> = enc.get_offsets().iter().filter(|(a, b)| a != b).collect();
    if let (Some(first), Some(last)) = (content.first(), content.last()) {
        for f in findings {
            f.boundary_fragment =
                (has_previous && f.start == first.0) || (has_next && f.end == last.1);
        }
    }
}
fn resolve(findings: Vec<Finding>) -> Vec<Finding> {
    let mut findings: Vec<Finding> = findings
        .iter()
        .filter(|f| {
            !(f.boundary_fragment
                && findings.iter().any(|g| {
                    g.label == f.label
                        && g.start <= f.start
                        && g.end >= f.end
                        && (g.start < f.start || g.end > f.end)
                }))
        })
        .cloned()
        .collect();
    findings.sort_by(|a, b| {
        b.raw_score
            .total_cmp(&a.raw_score)
            .then(a.start.cmp(&b.start))
            .then(a.end.cmp(&b.end))
            .then(a.label.cmp(&b.label))
    });
    let mut accepted: Vec<Finding> = Vec::new();
    for f in findings {
        if !accepted.iter().any(|a| f.start < a.end && a.start < f.end) {
            accepted.push(f);
        }
    }
    accepted.sort_by_key(|f| (f.start, f.end));
    accepted
}
fn redact(text: &str, findings: &[Finding]) -> Result<String> {
    let mut s = String::new();
    let mut pos = 0;
    for f in findings {
        ensure!(
            f.start >= pos
                && f.end <= text.len()
                && text.is_char_boundary(f.start)
                && text.is_char_boundary(f.end),
            "invalid UTF-8 finding boundary"
        );
        s.push_str(&text[pos..f.start]);
        s.push_str("[REDACTED]");
        pos = f.end;
    }
    s.push_str(&text[pos..]);
    Ok(s)
}
struct Detector {
    config: Config,
    tokenizer: Tokenizer,
    session: Session,
}
impl Detector {
    fn load(path: &Path) -> Result<Self> {
        let mut config: Config = serde_json::from_slice(&std::fs::read(path.join("config.json"))?)?;
        if let Some(file) = &config.calibration_file {
            config.calibration = Some(serde_json::from_slice(&std::fs::read(path.join(file))?)?);
        }
        ensure!(config.labels.len() == 7, "expected seven labels");
        ensure!(
            config.max_length > 2 && config.stride < config.max_length - 2,
            "invalid chunk settings"
        );
        if let Some(c) = &config.calibration {
            ensure!(
                !c.x.is_empty() && c.x.len() == c.y.len(),
                "invalid calibration lengths"
            );
            ensure!(
                c.x.iter().all(|x| x.is_finite())
                    && c.y.iter().all(|y| y.is_finite() && (0.0..=1.0).contains(y))
                    && c.x.windows(2).all(|v| v[0] < v[1])
                    && c.y.windows(2).all(|v| v[0] <= v[1]),
                "invalid calibration values"
            );
        }
        let mut tokenizer = Tokenizer::from_file(path.join("tokenizer.json"))
            .map_err(|e| anyhow::anyhow!("{e}"))?;
        tokenizer.with_padding(None);
        tokenizer
            .with_truncation(None)
            .map_err(|e| anyhow::anyhow!("{e}"))?;
        let _ = ort::init_from(path.join(if cfg!(target_os = "windows") { "onnxruntime.dll" } else if cfg!(target_os = "macos") { "libonnxruntime.dylib" } else { "libonnxruntime.so" }))?.commit();
        let session = Session::builder()?
            .with_intra_threads(1)
            .map_err(|e| anyhow::anyhow!("{e}"))?
            .commit_from_file(path.join("model.onnx"))?;
        Ok(Self {
            config,
            tokenizer,
            session,
        })
    }
    fn predict(&mut self, text: &str, debug: bool) -> Result<serde_json::Value> {
        let encodings = encode_windows(
            &self.tokenizer,
            text,
            self.config.max_length,
            self.config.stride,
        )?;
        let mut findings = Vec::new();
        let mut windows = Vec::new();
        for (index, e) in encodings.iter().enumerate() {
            let n = e.len();
            let ids: Vec<i64> = e.get_ids().iter().map(|x| *x as i64).collect();
            let mask: Vec<i64> = e.get_attention_mask().iter().map(|x| *x as i64).collect();
            let outputs=self.session.run(ort::inputs!["input_ids"=>Tensor::from_array(([1,n],ids.clone()))?,"attention_mask"=>Tensor::from_array(([1,n],mask.clone()))?])?;
            let (shape, logits) = outputs["logits"].try_extract_tensor::<f32>()?;
            ensure!(
                shape.as_ref() == [1, n as i64, 15],
                "unexpected logits shape: {shape:?}"
            );
            ensure!(logits.iter().all(|x| x.is_finite()), "non-finite logits");
            let mut decoded = decode(e, logits, &self.config.labels);
            mark_boundaries(e, &mut decoded, index > 0, index + 1 < encodings.len());
            if n == self.config.max_length {
                let content: Vec<_> = e.get_offsets().iter().filter(|(a, b)| a != b).collect();
                if let (Some(first), Some(last)) = (content.first(), content.last()) {
                    ensure!(
                        !decoded
                            .iter()
                            .any(|f| f.start == first.0 && f.end == last.1),
                        "entity spans a full window; complete entity boundary is unsupported"
                    );
                }
            }
            findings.extend(decoded);
            if debug {
                windows.push(json!({"input_ids":ids,"attention_mask":mask,"offsets":e.get_offsets(),"special_tokens_mask":e.get_special_tokens_mask(),"logits":logits.chunks_exact(15).collect::<Vec<_>>()}));
            }
        }
        let mut findings = resolve(findings);
        if let Some(c) = &self.config.calibration {
            for f in &mut findings {
                f.confidence = Some(c.apply(f.raw_score));
            }
            if let Some(t) = c.threshold {
                findings.retain(|f| f.confidence.unwrap() >= t);
            }
        }
        let redacted = redact(text, &findings)?;
        let mut result =
            json!({"findings":findings,"contains_pii":!findings.is_empty(),"redacted":redacted});
        if debug {
            result["windows"] = json!(windows);
        }
        Ok(result)
    }
}
fn main() -> Result<()> {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.is_empty() {
        bail!("usage: datafog-pii BUNDLE [--debug]; stdin JSONL objects with text:string");
    }
    ensure!(
        args.len() == 1 || (args.len() == 2 && args[1] == "--debug"),
        "unknown arguments"
    );
    let mut detector = Detector::load(Path::new(&args[0])).context("loading detector")?;
    let debug = args.len() == 2;
    let stdin = io::stdin();
    let mut stdout = io::BufWriter::new(io::stdout().lock());
    for line in stdin.lock().lines() {
        let line = line?;
        let result = (|| -> Result<_> {
            let input: serde_json::Value = serde_json::from_str(&line)?;
            let text = input
                .get("text")
                .and_then(|t| t.as_str())
                .context("expected text:string")?;
            detector.predict(text, debug)
        })();
        let output = match result {
            Ok(v) => v,
            Err(e) => json!({"error":format!("{e:#}")}),
        };
        writeln!(stdout, "{}", serde_json::to_string(&output)?)?;
        stdout.flush()?;
    }
    Ok(())
}
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn unicode_overlap_and_adjacent() {
        let f = |start, end, score| Finding {
            label: "first_name".into(),
            start,
            end,
            raw_score: score,
            confidence: None,
            boundary_fragment: false,
        };
        let spans = resolve(vec![f(5, 8, 0.6), f(5, 6, 0.9), f(6, 8, 0.8), f(5, 6, 0.7)]);
        assert_eq!(spans.len(), 2);
        assert_eq!(
            redact("🙂 José", &spans).unwrap(),
            "🙂 [REDACTED][REDACTED]é"
        );
        assert!(redact("🙂", &[f(1, 3, 1.)]).is_err());
    }
    #[test]
    fn oracle_boundaries_and_complete_token_coverage() {
        let fixture: serde_json::Value =
            serde_json::from_str(include_str!("../tests/boundary-oracle.json")).unwrap();
        let mut tokenizer = Tokenizer::from_file(
            Path::new(env!("CARGO_MANIFEST_DIR")).join("../bundle/tokenizer.json"),
        )
        .unwrap();
        tokenizer.with_truncation(None).unwrap();
        tokenizer.with_padding(None);
        let labels: Vec<String> = serde_json::from_value(fixture["labels"].clone()).unwrap();
        for case in fixture["cases"].as_array().unwrap() {
            let text = case["text"].as_str().unwrap();
            let windows = encode_windows(&tokenizer, text, 512, 42).unwrap();
            assert_eq!(windows.len(), case["windows"].as_array().unwrap().len());
            let mut findings = Vec::new();
            let mut covered = Vec::new();
            for (i, e) in windows.iter().enumerate() {
                let oracle = &case["windows"][i];
                assert_eq!(
                    serde_json::to_value(e.get_ids()).unwrap(),
                    oracle["input_ids"]
                );
                assert_eq!(
                    serde_json::to_value(e.get_offsets()).unwrap(),
                    oracle["offsets"]
                );
                covered.extend_from_slice(&e.get_ids()[if i == 0 { 1 } else { 43 }..e.len() - 1]);
                let logits: Vec<Vec<f32>> =
                    serde_json::from_value(oracle["logits"].clone()).unwrap();
                let mut f = decode(e, &logits.concat(), &labels);
                mark_boundaries(e, &mut f, i > 0, i + 1 < windows.len());
                findings.extend(f);
            }
            let full = tokenizer.encode(text, false).unwrap();
            assert_eq!(covered, full.get_ids());
            let actual = resolve(findings);
            let spans: Vec<_> = actual
                .iter()
                .map(|f| json!({"label":f.label,"start":f.start,"end":f.end}))
                .collect();
            assert_eq!(json!(spans), case["expected"]);
            assert_eq!(
                redact(text, &actual).unwrap(),
                case["redacted"].as_str().unwrap()
            );
        }
    }
    #[test]
    fn interpolation() {
        let c = Calibration {
            x: vec![0.2, 0.8],
            y: vec![0.1, 0.9],
            threshold: None,
        };
        assert!((c.apply(0.5) - 0.5).abs() < 1e-12);
        assert_eq!(c.apply(0.), 0.1);
        assert_eq!(c.apply(1.), 0.9);
    }
}