skymap-gen/src-tauri/src/main.rs

#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]

use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fs;
use std::io::{BufRead, BufReader, Read};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tauri::{async_runtime, Env, State, Window};

#[derive(Default, Clone)]
struct Paths {
    project_root: PathBuf,
    active_generation_pid: Arc<Mutex<Option<u32>>>,
    cancel_generation_requested: Arc<AtomicBool>,
}

#[derive(Serialize)]
struct MapInfo {
    path: String,
    filename: String,
    modified: i64,
}

#[derive(Serialize)]
struct GenerateResult {
    output_path: String,
}

#[derive(Debug, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
struct GenSettings {
    steps: Option<u32>,
    width: Option<u32>,
    height: Option<u32>,
    guidance: Option<f32>,
    scheduler: Option<String>,
    upscale: Option<String>,
    seam_inpaint: Option<bool>,
    model_path: Option<String>,
    base_model: Option<String>,
    vae_model: Option<String>,
}

fn discover_root() -> Result<PathBuf, String> {
    let cwd = std::env::current_dir().map_err(|e| format!("Failed to resolve current dir: {e}"))?;
    for ancestor in cwd.ancestors() {
        if ancestor.join("generate_equirect.py").exists() {
            return Ok(ancestor.to_path_buf());
        }
    }
    Err("Could not find project root (generate_equirect.py not found in ancestors)".into())
}

fn output_dir(root: &Path) -> Result<PathBuf, String> {
    let mut dir = root.to_path_buf();
    dir.push("output");
    fs::create_dir_all(&dir).map_err(|e| format!("Failed to create output dir: {e}"))?;
    Ok(dir)
}

fn script_path(root: &Path) -> Result<PathBuf, String> {
    let candidate = root.join("generate_equirect.py");
    if candidate.exists() {
        return Ok(candidate);
    }
    Err("generate_equirect.py not found".to_string())
}

/// Pick the Python binary to run: prefer project-local venv to ensure deps are present.
fn python_binary(root: &Path) -> PathBuf {
    let venv_py = root.join(".venv").join("bin").join("python");
    if venv_py.exists() {
        return venv_py;
    }
    PathBuf::from("python3")
}

#[tauri::command]
fn list_maps(state: State<Paths>) -> Result<Vec<MapInfo>, String> {
    let dir = output_dir(&state.project_root)?;
    let mut maps = Vec::new();
    if !dir.exists() {
        return Ok(maps);
    }
    for entry in fs::read_dir(&dir).map_err(|e| format!("Failed to read output dir: {e}"))? {
        let entry = entry.map_err(|e| format!("Failed to read entry: {e}"))?;
        let path = entry.path();
        if path
            .extension()
            .and_then(|s| s.to_str())
            .map(|s| s.eq_ignore_ascii_case("png"))
            != Some(true)
        {
            continue;
        }
        let metadata = entry
            .metadata()
            .map_err(|e| format!("Metadata error: {e}"))?;
        let modified = metadata
            .modified()
            .map(|m| {
                m.duration_since(std::time::UNIX_EPOCH)
                    .map(|d| d.as_secs() as i64)
                    .unwrap_or(0)
            })
            .unwrap_or(0);
        let filename = path
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or_default()
            .to_string();
        maps.push(MapInfo {
            path: path.to_string_lossy().to_string(),
            filename,
            modified,
        });
    }
    maps.sort_by(|a, b| b.modified.cmp(&a.modified));
    Ok(maps)
}

#[tauri::command]
fn delete_map(path: String, state: State<Paths>) -> Result<(), String> {
    let dir = output_dir(&state.project_root)?
        .canonicalize()
        .map_err(|e| format!("Failed to resolve output dir: {e}"))?;
    let candidate = PathBuf::from(path);
    let resolved = candidate
        .canonicalize()
        .map_err(|e| format!("Failed to resolve map path: {e}"))?;

    if !resolved.starts_with(&dir) {
        return Err("Refusing to delete a file outside the output directory".to_string());
    }
    if resolved
        .extension()
        .and_then(|s| s.to_str())
        .map(|s| s.eq_ignore_ascii_case("png"))
        != Some(true)
    {
        return Err("Only PNG maps can be deleted".to_string());
    }

    fs::remove_file(&resolved).map_err(|e| format!("Failed to delete map: {e}"))
}

fn terminate_process(pid: u32) -> Result<(), String> {
    #[cfg(target_family = "unix")]
    {
        Command::new("kill")
            .arg("-TERM")
            .arg(pid.to_string())
            .status()
            .map_err(|e| format!("Failed to request cancellation: {e}"))?;
        Ok(())
    }

    #[cfg(target_family = "windows")]
    {
        Command::new("taskkill")
            .args(["/PID", &pid.to_string(), "/T", "/F"])
            .status()
            .map_err(|e| format!("Failed to request cancellation: {e}"))?;
        Ok(())
    }
}

#[tauri::command]
fn cancel_generation(state: State<Paths>) -> Result<(), String> {
    state
        .cancel_generation_requested
        .store(true, Ordering::SeqCst);
    let pid = *state
        .active_generation_pid
        .lock()
        .map_err(|_| "Failed to read active generation state".to_string())?;

    match pid {
        Some(0) => Ok(()),
        Some(pid) => terminate_process(pid),
        None => Err("No generation is running".to_string()),
    }
}

#[tauri::command]
async fn generate_map(
    window: Window,
    prompt: String,
    settings: Option<GenSettings>,
    state: State<'_, Paths>,
) -> Result<GenerateResult, String> {
    if prompt.trim().is_empty() {
        return Err("Prompt must not be empty".into());
    }

    let prompt_clone = prompt.clone();
    let root = state.project_root.clone();
    let cfg = settings.unwrap_or_default();
    let steps = cfg.steps.unwrap_or(25);
    let width = cfg.width.unwrap_or(1536);
    let height = cfg.height.unwrap_or(768);
    let guidance = cfg.guidance.unwrap_or(6.5);
    let scheduler = cfg.scheduler.unwrap_or_else(|| "dpmsolver-sde".to_string());
    let upscale = cfg.upscale.unwrap_or_else(|| "none".to_string());
    let seam_inpaint = cfg.seam_inpaint.unwrap_or(false);
    let model_path = cfg
        .model_path
        .unwrap_or_else(|| "proximasan/sdxl-360-diffusion".to_string());
    let base_model = cfg
        .base_model
        .unwrap_or_else(|| "stabilityai/stable-diffusion-xl-base-1.0".to_string());
    let vae_model = cfg
        .vae_model
        .unwrap_or_else(|| "madebyollin/sdxl-vae-fp16-fix".to_string());
    let active_generation_pid = state.active_generation_pid.clone();
    let cancel_generation_requested = state.cancel_generation_requested.clone();
    let window = window.clone();

    {
        let mut active = active_generation_pid
            .lock()
            .map_err(|_| "Failed to update active generation state".to_string())?;
        if active.is_some() {
            return Err("Generation is already running".to_string());
        }
        *active = Some(0);
        cancel_generation_requested.store(false, Ordering::SeqCst);
    }

    async_runtime::spawn_blocking(move || {
        let result = (|| {
            let out_dir = output_dir(&root)?;
            let script = script_path(&root)?;
            let python = python_binary(&root);

            let mut cmd = Command::new(python);
            cmd.current_dir(&root)
                .arg(script.as_os_str())
                .arg("--prompt")
                .arg(prompt_clone)
                .arg("--output-dir")
                .arg(&out_dir)
                .arg("--work-dir")
                .arg(root.to_string_lossy().to_string())
                .arg("--upscale")
                .arg(upscale)
                .arg("--steps")
                .arg(steps.to_string())
                .arg("--guidance")
                .arg(guidance.to_string())
                .arg("--width")
                .arg(width.to_string())
                .arg("--height")
                .arg(height.to_string())
                .arg("--scheduler")
                .arg(scheduler)
                .arg("--model-path")
                .arg(model_path)
                .arg("--base-model")
                .arg(base_model)
                .arg("--vae-model")
                .arg(vae_model)
                .stdout(Stdio::piped())
                .stderr(Stdio::piped());

            if seam_inpaint {
                cmd.arg("--seam-inpaint");
            }

            let mut child = cmd
                .spawn()
                .map_err(|e| format!("Failed to start generator: {e}"))?;

            {
                let mut active = active_generation_pid
                    .lock()
                    .map_err(|_| "Failed to update active generation state".to_string())?;
                *active = Some(child.id());
            }

            if cancel_generation_requested.load(Ordering::SeqCst) {
                let _ = child.kill();
            }

            let stdout = child
                .stdout
                .take()
                .ok_or_else(|| "Failed to capture generator stdout".to_string())?;
            let stderr = child
                .stderr
                .take()
                .ok_or_else(|| "Failed to capture generator stderr".to_string())?;

            // Consume stderr in a side thread to avoid blocking.
            let mut stderr_reader = BufReader::new(stderr);
            let stderr_handle = std::thread::spawn(move || {
                let mut buf = String::new();
                let _ = stderr_reader.read_to_string(&mut buf);
                buf
            });

            let mut last_path: Option<String> = None;
            let mut reader = BufReader::new(stdout);
            let mut line = String::new();
            while reader
                .read_line(&mut line)
                .map_err(|e| format!("Failed to read generator output: {e}"))?
                > 0
            {
                let trimmed = line.trim_end().to_string();
                if let Some(json_str) = trimmed.strip_prefix("PROGRESS ") {
                    if let Ok(val) = serde_json::from_str::<Value>(json_str) {
                        let _ = window.emit("gen-progress", val);
                    }
                } else if !trimmed.is_empty() {
                    last_path = Some(trimmed.clone());
                }
                line.clear();
            }

            let status = child
                .wait()
                .map_err(|e| format!("Failed to wait for generator: {e}"))?;
            let stderr_output = stderr_handle.join().unwrap_or_default();

            if cancel_generation_requested.load(Ordering::SeqCst) {
                return Err("Generation cancelled".to_string());
            }

            if !status.success() {
                return Err(format!("Generator failed: {stderr_output}"));
            }

            let path_line =
                last_path.ok_or_else(|| "Generator did not return a path".to_string())?;
            let path = PathBuf::from(path_line.trim());
            let resolved = if path.is_absolute() {
                path
            } else {
                out_dir.join(path)
            };
            Ok(GenerateResult {
                output_path: resolved.to_string_lossy().to_string(),
            })
        })();

        if let Ok(mut active) = active_generation_pid.lock() {
            *active = None;
        }
        cancel_generation_requested.store(false, Ordering::SeqCst);

        result
    })
    .await
    .map_err(|e| format!("Generation task failed to join: {e}"))?
}

fn main() {
    let context = tauri::generate_context!();
    let env = Env::default();
    // Try to locate the project root by walking up from cwd; if missing, fall back to Tauri resource dir.
    let root = discover_root()
        .or_else(|_| {
            tauri::api::path::resource_dir(context.package_info(), &env)
                .ok_or_else(|| "Could not locate project root or resource dir".to_string())
        })
        .unwrap_or_else(|_| PathBuf::from("."));

    let state = Paths {
        project_root: root,
        ..Default::default()
    };
    tauri::Builder::default()
        .manage(state)
        .invoke_handler(tauri::generate_handler![
            list_maps,
            delete_map,
            cancel_generation,
            generate_map
        ])
        .run(context)
        .expect("error while running tauri application");
}