Adom Chipsmith

pub mod cache;
pub mod convert;

use anyhow::Result;
use rust_embed::RustEmbed;
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use tiny_http::{Header, Method, Request, Response, StatusCode};

#[derive(RustEmbed)]
#[folder = "vendor/3d-viewer-dist/"]
struct ViewerAssets;

const SHELL_HTML: &[u8] = include_bytes!("../assets/shell.html");
const ADOM_CSS: &[u8] = include_bytes!("../assets/adom.css");
const ICON_SVG: &[u8] = include_bytes!("../assets/icon.svg");

pub struct ServerConfig {
    pub port: u16,
    pub initial_glb_path: Option<PathBuf>,
    pub initial_source_label: Option<String>,
    /// Provenance: "kicad-library" | "local-file" | "upload" | "unknown".
    /// Drives the UI's "From KiCad library / Local file / Upload" badge so
    /// the user always knows where the geometry came from.
    pub initial_source_type: Option<String>,
}

#[derive(Clone, Debug)]
pub struct ConsoleMessage {
    pub seq: u64,
    pub ts_ms: u64,
    pub level: String,
    pub text: String,
}

#[derive(Clone, Debug)]
pub struct EvalJob {
    pub id: String,
    pub code: String,
    pub result: Option<serde_json::Value>,
    pub done: bool,
    pub taken: bool,
}

pub struct ServerState {
    pub config: ServerConfig,
    pub current_glb_path: Mutex<Option<PathBuf>>,
    pub current_source_label: Mutex<Option<String>>,
    pub current_source_type: Mutex<Option<String>>,
    pub current_materials: Mutex<serde_json::Value>,
    /// Last components tree the UI posted to /api/components.
    pub components: Mutex<serde_json::Value>,
    /// Per-node visibility map {nodeName: bool}.
    pub component_visibility: Mutex<serde_json::Value>,
    /// Active section planes — list of {axis, offset, enabled, flipped}.
    pub sections: Mutex<serde_json::Value>,
    /// Last selection posted by the UI (for `adom-step inspect`-like queries).
    pub selection: Mutex<serde_json::Value>,
    /// Pending command bus — UI polls /api/poll-command and applies.
    pub pending_command: Mutex<Option<serde_json::Value>>,
    /// JS eval pipeline.
    pub eval_jobs: Mutex<Vec<EvalJob>>,
    /// Console ring buffer (UI -> server -> CLI tail).
    pub console: Mutex<VecDeque<ConsoleMessage>>,
    pub console_seq: AtomicU64,
}

const CONSOLE_RING_CAP: usize = 2000;

impl ServerState {
    pub fn new(config: ServerConfig) -> Arc<Self> {
        let initial_path = config.initial_glb_path.clone();
        let initial_label = config.initial_source_label.clone();
        let initial_type = config.initial_source_type.clone();
        Arc::new(Self {
            config,
            current_glb_path: Mutex::new(initial_path),
            current_source_label: Mutex::new(initial_label),
            current_source_type: Mutex::new(initial_type),
            current_materials: Mutex::new(serde_json::json!([])),
            components: Mutex::new(serde_json::json!([])),
            component_visibility: Mutex::new(serde_json::json!({})),
            sections: Mutex::new(serde_json::json!([])),
            selection: Mutex::new(serde_json::json!(null)),
            pending_command: Mutex::new(None),
            eval_jobs: Mutex::new(Vec::new()),
            console: Mutex::new(VecDeque::with_capacity(CONSOLE_RING_CAP)),
            console_seq: AtomicU64::new(0),
        })
    }

    pub fn push_console(&self, level: &str, text: &str) {
        let ts_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_millis() as u64)
            .unwrap_or(0);
        let seq = self.console_seq.fetch_add(1, Ordering::SeqCst);
        let mut buf = self.console.lock().unwrap();
        if buf.len() >= CONSOLE_RING_CAP {
            buf.pop_front();
        }
        buf.push_back(ConsoleMessage {
            seq,
            ts_ms,
            level: level.to_string(),
            text: text.to_string(),
        });
    }
}

pub fn run(state: Arc<ServerState>) -> Result<()> {
    let addr = format!("127.0.0.1:{}", state.config.port);
    let server = tiny_http::Server::http(&addr)
        .map_err(|e| anyhow::anyhow!("failed to bind {}: {}", addr, e))?;

    eprintln!("adom-step listening on http://{}", addr);

    for request in server.incoming_requests() {
        let state = state.clone();
        std::thread::spawn(move || {
            if let Err(e) = route(&state, request) {
                eprintln!("request error: {e:#}");
            }
        });
    }
    Ok(())
}

// Cloudflare-compliant anti-cache headers (per gallia/skills/adom guides).
const CACHE_CONTROL_VAL: &str = "no-store, no-cache, must-revalidate, max-age=0";
const PRAGMA_VAL: &str = "no-cache";
const SURROGATE_VAL: &str = "no-store";

fn no_cache<R: std::io::Read>(resp: Response<R>) -> Response<R> {
    resp.with_header(Header::from_bytes(&b"Cache-Control"[..], CACHE_CONTROL_VAL.as_bytes()).unwrap())
        .with_header(Header::from_bytes(&b"Pragma"[..], PRAGMA_VAL.as_bytes()).unwrap())
        .with_header(Header::from_bytes(&b"Surrogate-Control"[..], SURROGATE_VAL.as_bytes()).unwrap())
}

fn json_response(status: u16, body: &serde_json::Value) -> Response<std::io::Cursor<Vec<u8>>> {
    let bytes = serde_json::to_vec(body).unwrap_or_default();
    let resp = Response::from_data(bytes)
        .with_status_code(StatusCode(status))
        .with_header(Header::from_bytes(&b"Content-Type"[..], &b"application/json"[..]).unwrap());
    no_cache(resp)
}

fn text_response(status: u16, mime: &str, body: Vec<u8>) -> Response<std::io::Cursor<Vec<u8>>> {
    let resp = Response::from_data(body)
        .with_status_code(StatusCode(status))
        .with_header(Header::from_bytes(&b"Content-Type"[..], mime.as_bytes()).unwrap());
    no_cache(resp)
}

fn read_body(req: &mut Request) -> Vec<u8> {
    let mut buf = Vec::new();
    let _ = req.as_reader().read_to_end(&mut buf);
    buf
}

fn parse_json(req: &mut Request) -> Option<serde_json::Value> {
    let bytes = read_body(req);
    serde_json::from_slice(&bytes).ok()
}

fn next_eval_id(state: &ServerState) -> String {
    let seq = state.console_seq.fetch_add(1, Ordering::SeqCst);
    let ts = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0);
    format!("e_{ts}_{seq}")
}

fn route(state: &Arc<ServerState>, mut request: Request) -> Result<()> {
    let url = request.url().to_string();
    // Strip query string for routing — handlers re-parse from request.url() if needed.
    let path = url.split('?').next().unwrap_or(&url).to_string();
    let method = request.method().clone();

    // ---- Health & state ---------------------------------------------------

    if path == "/health" {
        let body = serde_json::json!({
            "ok": true,
            "service": "adom-step",
            "port": state.config.port,
        });
        return Ok(request.respond(json_response(200, &body))?);
    }

    if path == "/state" {
        let body = serde_json::json!({
            "source": *state.current_source_label.lock().unwrap(),
            "source_type": *state.current_source_type.lock().unwrap(),
            "has_model": state.current_glb_path.lock().unwrap().is_some(),
            "materials": *state.current_materials.lock().unwrap(),
            "components": *state.components.lock().unwrap(),
            "visibility": *state.component_visibility.lock().unwrap(),
            "sections": *state.sections.lock().unwrap(),
        });
        return Ok(request.respond(json_response(200, &body))?);
    }

    if path == "/shutdown" && method == Method::Post {
        let _ = request.respond(json_response(200, &serde_json::json!({"ok": true})));
        std::process::exit(0);
    }

    // ---- Static shell -----------------------------------------------------

    if path == "/" || path == "/index.html" {
        return Ok(request.respond(text_response(200, "text/html; charset=utf-8", SHELL_HTML.to_vec()))?);
    }
    if path == "/adom.css" {
        return Ok(request.respond(text_response(200, "text/css; charset=utf-8", ADOM_CSS.to_vec()))?);
    }
    if path == "/favicon.svg" || path == "/favicon.ico" {
        return Ok(request.respond(text_response(200, "image/svg+xml", ICON_SVG.to_vec()))?);
    }

    // Vendored Babylon9 ESM bundle and chunks: /viewer/<filename>
    if let Some(asset) = path.strip_prefix("/viewer/") {
        if let Some(file) = ViewerAssets::get(asset) {
            let mime = if asset.ends_with(".js") {
                "application/javascript; charset=utf-8"
            } else if asset.ends_with(".css") {
                "text/css; charset=utf-8"
            } else if asset.ends_with(".env") {
                "application/octet-stream"
            } else {
                "application/octet-stream"
            };
            return Ok(request.respond(text_response(200, mime, file.data.into_owned()))?);
        }
        return Ok(request.respond(json_response(404, &serde_json::json!({"error": "not found"})))?);
    }

    // ---- Model serving ----------------------------------------------------

    if path == "/model.glb" {
        let glb_path = state.current_glb_path.lock().unwrap().clone();
        match glb_path {
            Some(p) => match std::fs::read(&p) {
                Ok(bytes) => {
                    let resp = Response::from_data(bytes)
                        .with_status_code(StatusCode(200))
                        .with_header(Header::from_bytes(&b"Content-Type"[..], &b"model/gltf-binary"[..]).unwrap());
                    return Ok(request.respond(no_cache(resp))?);
                }
                Err(e) => {
                    return Ok(request.respond(json_response(
                        500,
                        &serde_json::json!({"error": format!("read {}: {e}", p.display())}),
                    ))?);
                }
            },
            None => {
                return Ok(request.respond(json_response(
                    404,
                    &serde_json::json!({"error": "no model loaded — POST a STEP file to /api/load or drop one onto the canvas"}),
                ))?);
            }
        }
    }

    // ---- STEP upload (drop zone + AI) ------------------------------------

    if path == "/api/load" && method == Method::Post {
        let bytes = read_body(&mut request);
        if bytes.is_empty() {
            return Ok(request.respond(json_response(
                400,
                &serde_json::json!({"error": "empty body"}),
            ))?);
        }
        // Filename hint via X-Filename header (for the source label in HUD).
        let filename = request
            .headers()
            .iter()
            .find(|h| h.field.as_str().as_str().eq_ignore_ascii_case("X-Filename"))
            .map(|h| h.value.as_str().to_string())
            .unwrap_or_else(|| "uploaded.step".to_string());

        match convert::convert_step_bytes(&bytes) {
            Ok(result) => {
                *state.current_glb_path.lock().unwrap() = Some(result.glb_path.clone());
                *state.current_source_label.lock().unwrap() = Some(filename.clone());
                *state.current_source_type.lock().unwrap() = Some("upload".to_string());
                *state.current_materials.lock().unwrap() = result.materials.clone();
                let body = serde_json::json!({
                    "ok": true,
                    "source": filename,
                    "glb_size_bytes": result.size_bytes,
                    "meshes": result.meshes,
                    "nodes": result.nodes,
                    "duration_ms": result.duration_ms,
                    "materials": result.materials,
                    "cache_hit": result.cache_hit,
                });
                return Ok(request.respond(json_response(200, &body))?);
            }
            Err(e) => {
                return Ok(request.respond(json_response(
                    500,
                    &serde_json::json!({"error": e.to_string()}),
                ))?);
            }
        }
    }

    // ---- Components tree (UI <-> server) ---------------------------------

    if path == "/api/components" && method == Method::Post {
        if let Some(v) = parse_json(&mut request) {
            *state.components.lock().unwrap() = v;
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }
    if path == "/api/components" && method == Method::Get {
        let v = state.components.lock().unwrap().clone();
        return Ok(request.respond(json_response(200, &v))?);
    }

    // ---- Visibility ------------------------------------------------------

    if path == "/api/visibility" && method == Method::Post {
        if let Some(v) = parse_json(&mut request) {
            *state.component_visibility.lock().unwrap() = v;
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }

    // ---- Sections --------------------------------------------------------

    if path == "/api/sections" && method == Method::Post {
        if let Some(v) = parse_json(&mut request) {
            *state.sections.lock().unwrap() = v;
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }
    if path == "/api/sections" && method == Method::Get {
        let v = state.sections.lock().unwrap().clone();
        return Ok(request.respond(json_response(200, &v))?);
    }

    // ---- Selection (UI -> server, drives `inspect` CLI) ------------------

    if path == "/api/selection" && method == Method::Post {
        if let Some(v) = parse_json(&mut request) {
            *state.selection.lock().unwrap() = v;
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }
    if path == "/api/selection" && method == Method::Get {
        let v = state.selection.lock().unwrap().clone();
        return Ok(request.respond(json_response(200, &v))?);
    }

    // ---- Command bus (CLI -> server -> UI poll) --------------------------

    if path == "/api/command" && method == Method::Post {
        if let Some(v) = parse_json(&mut request) {
            *state.pending_command.lock().unwrap() = Some(v);
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }
    if path == "/api/poll-command" && method == Method::Get {
        let v = state.pending_command.lock().unwrap().take();
        return Ok(request.respond(json_response(200, &serde_json::json!({"command": v})))?);
    }

    // ---- JS eval pipeline ------------------------------------------------

    if path == "/eval" && method == Method::Post {
        let body: serde_json::Value = parse_json(&mut request).unwrap_or(serde_json::json!({}));
        let code = body
            .get("code")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();
        if code.is_empty() {
            return Ok(request.respond(json_response(
                400,
                &serde_json::json!({"error": "missing 'code'"}),
            ))?);
        }
        let id = next_eval_id(state);
        state.eval_jobs.lock().unwrap().push(EvalJob {
            id: id.clone(),
            code,
            result: None,
            done: false,
            taken: false,
        });
        return Ok(request.respond(json_response(200, &serde_json::json!({"id": id})))?);
    }

    if path == "/eval/pending" && method == Method::Get {
        let mut jobs = state.eval_jobs.lock().unwrap();
        let next = jobs
            .iter_mut()
            .find(|j| !j.done && !j.taken)
            .map(|j| {
                j.taken = true;
                serde_json::json!({"id": j.id, "code": j.code})
            });
        return Ok(request.respond(json_response(200, &serde_json::json!({"job": next})))?);
    }

    if path.starts_with("/eval/") && path.ends_with("/result") && method == Method::Post {
        let id = path
            .trim_start_matches("/eval/")
            .trim_end_matches("/result")
            .to_string();
        let body = parse_json(&mut request).unwrap_or(serde_json::json!(null));
        let mut jobs = state.eval_jobs.lock().unwrap();
        if let Some(job) = jobs.iter_mut().find(|j| j.id == id) {
            job.result = Some(body);
            job.done = true;
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }

    if path.starts_with("/eval/")
        && !path.ends_with("/result")
        && path != "/eval/pending"
        && method == Method::Get
    {
        let id = path.trim_start_matches("/eval/").to_string();
        let jobs = state.eval_jobs.lock().unwrap();
        let body = match jobs.iter().find(|j| j.id == id) {
            Some(j) => serde_json::json!({
                "id": j.id,
                "done": j.done,
                "result": j.result.clone().unwrap_or(serde_json::Value::Null),
            }),
            None => serde_json::json!({"error": "no such id"}),
        };
        return Ok(request.respond(json_response(200, &body))?);
    }

    // ---- Console forwarder ----------------------------------------------

    if path == "/console" && method == Method::Post {
        if let Some(v) = parse_json(&mut request) {
            let level = v.get("level").and_then(|x| x.as_str()).unwrap_or("log");
            let text = v.get("text").and_then(|x| x.as_str()).unwrap_or("");
            state.push_console(level, text);
        }
        return Ok(request.respond(json_response(200, &serde_json::json!({"ok": true})))?);
    }

    if path == "/console" && method == Method::Get {
        let qs = url.split('?').nth(1).unwrap_or("");
        let since: u64 = qs
            .split('&')
            .find_map(|p| p.strip_prefix("since=").and_then(|v| v.parse().ok()))
            .unwrap_or(0);
        let buf = state.console.lock().unwrap();
        let messages: Vec<_> = buf
            .iter()
            .filter(|m| m.seq >= since)
            .map(|m| {
                serde_json::json!({
                    "seq": m.seq,
                    "ts_ms": m.ts_ms,
                    "level": m.level,
                    "text": m.text,
                })
            })
            .collect();
        let next_since = buf.iter().last().map(|m| m.seq + 1).unwrap_or(since);
        let body = serde_json::json!({"messages": messages, "next_since": next_since});
        return Ok(request.respond(json_response(200, &body))?);
    }

    // ---- 404 -------------------------------------------------------------

    Ok(request.respond(json_response(
        404,
        &serde_json::json!({"error": format!("no route for {} {}", method, path)}),
    ))?)
}