AI Flow: Reverse Engineering

//! KiCad's own picture of a board: the shared service plots gerbers and an Excellon drill file, and this
//! module draws them into board-frame pixels at 40 px/mm. It reads only what KiCad writes (RS-274X with
//! apertures, aperture macros, regions and arcs, dark polarity), so what `check fidelity` scores is what
//! KiCad draws, not what the flow meant to write.
use crate::raster::Mask;
use std::collections::HashMap;
use std::path::Path;

/// A raster of the board frame: px = (mm - origin) * res, y down, pixel (u, v) centred on (u + 0.5, v + 0.5).
#[derive(Clone, Copy)]
pub struct Frame { pub w: usize, pub h: usize, pub res: f64, pub ox: f64, pub oy: f64 }
impl Frame {
    pub fn px(&self, x: f64, y: f64) -> (f64, f64) { ((x - self.ox) * self.res, (y - self.oy) * self.res) }
}

pub fn fill_poly(m: &mut Mask, f: &Frame, pts_mm: &[(f64, f64)]) {
    if pts_mm.len() < 3 { return; }
    let p: Vec<(f64, f64)> = pts_mm.iter().map(|q| f.px(q.0, q.1)).collect();
    let (y0, y1) = p.iter().fold((f64::MAX, f64::MIN), |a, q| (a.0.min(q.1), a.1.max(q.1)));
    let v0 = (y0 - 0.5).ceil().max(0.0) as usize; let v1 = ((y1 - 0.5).floor().min(f.h as f64 - 1.0)) as i64;
    if v1 < v0 as i64 { return; }
    let mut rows: Vec<Vec<f64>> = vec![vec![]; v1 as usize - v0 + 1];
    for i in 0..p.len() {
        let (a, b) = (p[i], p[(i + 1) % p.len()]);
        if a.1 == b.1 { continue; }
        let (lo, hi) = if a.1 < b.1 { (a, b) } else { (b, a) };
        let s = ((lo.1 - 0.5).ceil().max(v0 as f64)) as usize; let e = (hi.1 - 0.5).ceil().min(v1 as f64 + 1.0) as usize;
        for v in s..e { let yc = v as f64 + 0.5; rows[v - v0].push(lo.0 + (yc - lo.1) / (hi.1 - lo.1) * (hi.0 - lo.0)); }
    }
    for (k, xs) in rows.iter_mut().enumerate() {
        xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let v = v0 + k;
        for c in xs.chunks(2) { if c.len() < 2 { continue; }
            let (u0, u1) = ((c[0] - 0.5).ceil().max(0.0) as usize, (c[1] - 0.5).floor().min(f.w as f64 - 1.0) as i64);
            if u1 < u0 as i64 { continue; }
            for u in u0..=u1 as usize { m.d[v * f.w + u] = true; } }
    }
}

pub fn fill_disk(m: &mut Mask, f: &Frame, x: f64, y: f64, r_mm: f64) {
    let (cx, cy) = f.px(x, y); let r = r_mm * f.res;
    let (v0, v1) = (((cy - r - 1.0).floor().max(0.0)) as usize, ((cy + r + 1.0).ceil().min(f.h as f64)) as usize);
    let (u0, u1) = (((cx - r - 1.0).floor().max(0.0)) as usize, ((cx + r + 1.0).ceil().min(f.w as f64)) as usize);
    for v in v0..v1 { for u in u0..u1 { let (dx, dy) = (u as f64 + 0.5 - cx, v as f64 + 0.5 - cy); if dx * dx + dy * dy <= r * r { m.d[v * f.w + u] = true; } } }
}

/// A thick line with round ends.
pub fn fill_seg(m: &mut Mask, f: &Frame, a: (f64, f64), b: (f64, f64), wd: f64) {
    fill_disk(m, f, a.0, a.1, wd / 2.0); fill_disk(m, f, b.0, b.1, wd / 2.0);
    let (dx, dy) = (b.0 - a.0, b.1 - a.1); let l = (dx * dx + dy * dy).sqrt(); if l < 1e-9 { return; }
    let (nx, ny) = (-dy / l * wd / 2.0, dx / l * wd / 2.0);
    fill_poly(m, f, &[(a.0 + nx, a.1 + ny), (b.0 + nx, b.1 + ny), (b.0 - nx, b.1 - ny), (a.0 - nx, a.1 - ny)]);
}

#[derive(Clone, Debug)]
enum Prim { Disk(f64, f64, f64), Poly(Vec<(f64, f64)>), Seg((f64, f64), (f64, f64), f64) }

fn rot(p: (f64, f64), deg: f64) -> (f64, f64) { let (c, s) = (deg.to_radians().cos(), deg.to_radians().sin()); (p.0 * c - p.1 * s, p.0 * s + p.1 * c) }

/// $n and sums like $1+$1 in a macro expression.
fn eval(e: &str, a: &[f64]) -> f64 {
    let term = |t: &str| -> f64 { let t = t.trim(); if let Some(n) = t.strip_prefix('$') { a.get(n.parse::<usize>().unwrap_or(1).wrapping_sub(1)).copied().unwrap_or(0.0) } else { t.parse().unwrap_or(0.0) } };
    e.split('+').map(|s| s.split('x').map(term).product::<f64>()).sum()
}

fn macro_prims(body: &[String], args: &[f64]) -> Vec<Prim> {
    let mut out = vec![];
    for line in body {
        let f: Vec<&str> = line.split(',').collect();
        let num = |i: usize| eval(f.get(i).copied().unwrap_or("0"), args);
        match f[0].trim() {
            "1" => out.push(Prim::Disk(num(3), num(4), num(2) / 2.0)),
            "4" => { let n = num(2) as usize; let ang = num(5 + 2 * n); out.push(Prim::Poly((0..=n).map(|k| rot((num(3 + 2 * k), num(4 + 2 * k)), ang)).collect())); }
            "20" => { let ang = num(7); out.push(Prim::Seg(rot((num(3), num(4)), ang), rot((num(5), num(6)), ang), num(2))); }
            "21" => { let (w, h, cx, cy, ang) = (num(2), num(3), num(4), num(5), num(6)); out.push(Prim::Poly([(-w / 2.0, -h / 2.0), (w / 2.0, -h / 2.0), (w / 2.0, h / 2.0), (-w / 2.0, h / 2.0)].iter().map(|q| rot((q.0 + cx, q.1 + cy), ang)).collect())); }
            _ => {}
        }
    }
    out
}

fn draw(m: &mut Mask, f: &Frame, pr: &[Prim], at: (f64, f64)) {
    // gerber y runs up: the board's y = -y
    let w = |p: (f64, f64)| (at.0 + p.0, -(at.1 + p.1));
    for q in pr { match q {
        Prim::Disk(x, y, r) => { let c = w((*x, *y)); fill_disk(m, f, c.0, c.1, *r); }
        Prim::Poly(v) => { let pts: Vec<(f64, f64)> = v.iter().map(|p| w(*p)).collect(); fill_poly(m, f, &pts); }
        Prim::Seg(a, b, wd) => fill_seg(m, f, w(*a), w(*b), *wd),
    } }
}

fn arc_points(a: (f64, f64), b: (f64, f64), c: (f64, f64), cw: bool) -> Vec<(f64, f64)> {
    let r = ((a.0 - c.0).powi(2) + (a.1 - c.1).powi(2)).sqrt();
    let (t0, mut t1) = ((a.1 - c.1).atan2(a.0 - c.0), (b.1 - c.1).atan2(b.0 - c.0));
    if cw { if t1 >= t0 { t1 -= std::f64::consts::TAU; } } else if t1 <= t0 { t1 += std::f64::consts::TAU; }
    let n = (((t1 - t0).abs() * r * 40.0) as usize).clamp(4, 400);
    (1..=n).map(|k| { let t = t0 + (t1 - t0) * k as f64 / n as f64; (c.0 + r * t.cos(), c.1 + r * t.sin()) }).collect()
}

/// Draw one RS-274X file. Mask files plot their openings, so the result is "plotted", whatever the polarity tag says.
pub fn rasterize(text: &str, f: &Frame) -> Result<Mask, String> {
    let mut m = Mask::new(f.w, f.h);
    let mut macros: HashMap<String, Vec<String>> = HashMap::new();
    let mut aps: HashMap<i64, Vec<Prim>> = HashMap::new();
    let mut ap_w: HashMap<i64, f64> = HashMap::new();
    let (mut cur, mut pos, mut interp, mut region): (i64, (f64, f64), u8, Option<Vec<(f64, f64)>>) = (0, (0.0, 0.0), 1, None);
    let (mut xd, mut yd) = (6i32, 6i32);
    // split into commands: %...% blocks stay whole, the rest splits on '*'
    let mut rest = text;
    while !rest.is_empty() {
        let t = rest.trim_start();
        if t.is_empty() { break; }
        if let Some(b) = t.strip_prefix('%') {
            let end = b.find('%').ok_or("unterminated % block")?;
            let block = &b[..end]; rest = &b[end + 1..];
            let bt = block.trim();
            if let Some(fs) = bt.strip_prefix("FSLA") {
                let g = |c: char| fs.find(c).and_then(|i| fs[i + 2..].chars().next()).and_then(|d| d.to_digit(10)).unwrap_or(6) as i32;
                xd = g('X'); yd = g('Y');
            } else if bt.starts_with("MOIN") { return Err("inch gerbers are not expected from KiCad here".into()); }
            else if let Some(am) = bt.strip_prefix("AM") {
                let mut lines = am.split('*').map(|s| s.trim().replace('\n', "")).filter(|s| !s.is_empty());
                let name = lines.next().unwrap_or_default();
                macros.insert(name, lines.filter(|l| !l.starts_with('0')).collect());
            } else if let Some(ad) = bt.strip_prefix("ADD") {
                let ad = ad.trim_end_matches('*');
                let dig = ad.find(|c: char| !c.is_ascii_digit()).unwrap_or(ad.len());
                let code: i64 = ad[..dig].parse().map_err(|_| format!("bad aperture {ad}"))?;
                let (name, args) = ad[dig..].split_once(',').unwrap_or((&ad[dig..], ""));
                let nums: Vec<f64> = args.split('X').filter_map(|s| s.trim().parse().ok()).collect();
                let pr = match name {
                    "C" => vec![Prim::Disk(0.0, 0.0, nums.first().copied().unwrap_or(0.0) / 2.0)],
                    "R" => { let (w, h) = (nums[0] / 2.0, nums.get(1).copied().unwrap_or(nums[0]) / 2.0); vec![Prim::Poly(vec![(-w, -h), (w, -h), (w, h), (-w, h)])] }
                    "O" => { let (w, h) = (nums[0], nums.get(1).copied().unwrap_or(nums[0])); let r = w.min(h) / 2.0;
                             if w >= h { vec![Prim::Seg((-w / 2.0 + r, 0.0), (w / 2.0 - r, 0.0), 2.0 * r)] } else { vec![Prim::Seg((0.0, -h / 2.0 + r), (0.0, h / 2.0 - r), 2.0 * r)] } }
                    other => { let b = macros.get(other).ok_or(format!("unknown macro {other}"))?;
                               // KiCad's FreePoly apertures take a rotation as their one argument
                               let mut p = macro_prims(b, &nums);
                               if other.starts_with("FreePoly") { let ang = nums.first().copied().unwrap_or(0.0); p = macro_prims(b, &[]).into_iter().map(|q| match q { Prim::Poly(v) => Prim::Poly(v.into_iter().map(|x| rot(x, ang)).collect()), Prim::Disk(x, y, r) => { let c = rot((x, y), ang); Prim::Disk(c.0, c.1, r) } Prim::Seg(a, b, w) => Prim::Seg(rot(a, ang), rot(b, ang), w) }).collect(); }
                               p }
                };
                if let Some(Prim::Disk(_, _, r)) = pr.first() { ap_w.insert(code, 2.0 * *r); }
                aps.insert(code, pr);
            }
            continue;
        }
        let end = t.find('*').ok_or("unterminated command")?;
        let cmd = t[..end].trim(); rest = &t[end + 1..];
        if cmd.starts_with("G04") || cmd.is_empty() { continue; }
        if cmd == "M02" { break; }
        match cmd { "G01" => { interp = 1; continue; } "G02" => { interp = 2; continue; } "G03" => { interp = 3; continue; } "G75" | "G74" => continue,
            "G36" => { region = Some(vec![]); continue; }
            "G37" => { if let Some(r) = region.take() { fill_poly(&mut m, f, &r); } continue; } _ => {} }
        if let Some(d) = cmd.strip_prefix('D') { if let Ok(n) = d.parse::<i64>() { if n >= 10 { cur = n; continue; } } }
        if cmd.starts_with("G54D") { cur = cmd[4..].parse().unwrap_or(cur); continue; }
        // a coordinate op: optional G0x prefix, X, Y, I, J, Dnn
        let mut s = cmd;
        if s.starts_with("G01") { interp = 1; s = &s[3..]; } else if s.starts_with("G02") { interp = 2; s = &s[3..]; } else if s.starts_with("G03") { interp = 3; s = &s[3..]; }
        let (mut nx, mut ny, mut ci, mut cj, mut op) = (pos.0, pos.1, 0.0, 0.0, 0i64);
        let b = s.as_bytes(); let mut i = 0;
        while i < b.len() {
            let c = b[i] as char; let st = i + 1; let mut e = st;
            while e < b.len() && ((b[e] as char).is_ascii_digit() || b[e] == b'-' || b[e] == b'+') { e += 1; }
            let v: f64 = s[st..e].parse().unwrap_or(0.0);
            match c { 'X' => nx = v / 10f64.powi(xd), 'Y' => ny = v / 10f64.powi(yd), 'I' => ci = v / 10f64.powi(xd), 'J' => cj = v / 10f64.powi(yd), 'D' => op = v as i64, _ => {} }
            i = e;
        }
        let to = (nx, ny);
        match op {
            1 => {
                let pts: Vec<(f64, f64)> = if interp == 1 { vec![to] } else { arc_points(pos, to, (pos.0 + ci, pos.1 + cj), interp == 2) };
                if let Some(r) = region.as_mut() { if r.is_empty() { r.push((pos.0, -pos.1)); } for p in &pts { r.push((p.0, -p.1)); } }
                else { let w = ap_w.get(&cur).copied().unwrap_or(0.05); let mut a = pos; for p in &pts { fill_seg(&mut m, f, (a.0, -a.1), (p.0, -p.1), w); a = *p; } }
            }
            2 => { if let Some(r) = region.as_mut() { if !r.is_empty() { let done = std::mem::take(r); fill_poly(&mut m, f, &done); } } }
            3 => { if let Some(pr) = aps.get(&cur) { draw(&mut m, f, pr, to); } }
            _ => {}
        }
        pos = to;
    }
    Ok(m)
}

#[derive(Clone, Debug)]
pub struct Hole { pub x: f64, pub y: f64, pub d: f64, pub via: bool, pub plated: bool }

/// KiCad's Excellon file: tools, their function comments (via or not) and the hits (slots as their ends).
pub fn parse_drill(text: &str) -> Vec<Hole> {
    let (mut tools, mut out): (HashMap<String, (f64, bool, bool)>, Vec<Hole>) = (HashMap::new(), vec![]);
    let (mut cur, mut pend_via, mut pend_plated) = (String::new(), false, true);
    let mut inch = false;
    for line in text.lines().map(str::trim) {
        if line.starts_with("; #@! TA.AperFunction") { pend_via = line.contains("ViaDrill"); pend_plated = !line.contains("NonPlated"); continue; }
        if line == "INCH" || line.starts_with("INCH,") { inch = true; continue; }
        if line.starts_with('T') {
            if let Some(cpos) = line.find('C') { let d: f64 = line[cpos + 1..].parse().unwrap_or(0.0); tools.insert(line[..cpos].to_string(), (if inch { d * 25.4 } else { d }, pend_via, pend_plated)); }
            else if line != "T0" { cur = line.to_string(); }
            continue;
        }
        if line.starts_with('X') {
            let body = line.split("G85").next().unwrap_or(line);
            let y_at = body.find('Y').unwrap_or(body.len());
            let x: f64 = body[1..y_at].parse().unwrap_or(0.0); let y: f64 = body.get(y_at + 1..).and_then(|s| s.parse().ok()).unwrap_or(0.0);
            let k = if inch { 25.4 } else { 1.0 };
            if let Some(t) = tools.get(&cur) { out.push(Hole { x: x * k, y: -y * k, d: t.0, via: t.1, plated: t.2 }); }
        }
    }
    out
}

pub struct Plot { pub layers: HashMap<String, String>, pub drill: Vec<Hole> }

/// Ask the shared KiCad service for the board's plot. Fails plainly if the service cannot read the board.
pub fn plot(pcb: &Path, out_dir: &Path) -> Result<Plot, String> {
    let _ = std::fs::remove_dir_all(out_dir); std::fs::create_dir_all(out_dir.parent().unwrap_or(Path::new("."))).ok();
    let o = std::process::Command::new("service-kicad").args(["pcb", "export", "gerbers"]).arg(pcb).arg("--out").arg(out_dir).output().map_err(|e| format!("service-kicad not runnable: {e}"))?;
    if !o.status.success() { return Err(format!("KiCad could not plot the board: {}", String::from_utf8_lossy(&o.stderr).chars().take(300).collect::<String>())); }
    let mut layers = HashMap::new(); let mut drill = vec![];
    for e in std::fs::read_dir(out_dir).map_err(|e| e.to_string())? {
        let e = e.map_err(|e| e.to_string())?; let n = e.file_name().to_string_lossy().to_string();
        let txt = std::fs::read_to_string(e.path()).unwrap_or_default();
        if n.ends_with(".drl") { drill = parse_drill(&txt); continue; }
        for key in ["F_Cu", "B_Cu", "F_Mask", "B_Mask", "F_Silkscreen", "B_Silkscreen", "Edge_Cuts"] { if n.contains(key) { layers.insert(key.to_string(), txt.clone()); } }
    }
    if !layers.contains_key("F_Cu") { return Err("the plot has no F_Cu file".into()); }
    Ok(Plot { layers, drill })
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn draws_flashes_regions_and_macros() {
        let f = Frame { w: 200, h: 200, res: 40.0, ox: 100.0, oy: 100.0 };
        let g = "%FSLAX46Y46*%\n%MOMM*%\n%ADD10R,1.000000X0.500000*%\n%ADD11C,0.500000*%\nD10*\nX101000000Y-101000000D03*\nD11*\nG36*\nX102000000Y-102000000D02*\nG01*\nX103000000Y-102000000D01*\nX103000000Y-103000000D01*\nX102000000Y-103000000D01*\nX102000000Y-102000000D01*\nG37*\nM02*\n";
        let m = rasterize(g, &f).unwrap();
        assert!((m.count() as f64 - 40.0 * 20.0 - 1600.0).abs() < 90.0, "{}", m.count());
        assert!(m.at(40, 40) && m.at(100, 100) && !m.at(70, 70));
        let h = parse_drill("M48\n; #@! TA.AperFunction,Plated,PTH,ViaDrill\nT1C0.300\n; #@! TA.AperFunction,Plated,PTH,ComponentDrill\nT2C1.000\n%\nT1\nX101.0Y-102.0\nT2\nX103.5Y-104.0\nM30\n");
        assert_eq!(h.len(), 2); assert!(h[0].via && !h[1].via); assert!((h[1].y - 104.0).abs() < 1e-9 && (h[1].d - 1.0).abs() < 1e-9);
    }
}