swarm-engine-ui 0.1.6

//! LoRA Fine-tuning commands

use std::io::{BufRead, Write};
use std::path::PathBuf;

use clap::Subcommand;
use swarm_engine_core::learn::{ActionRecord, LearnModel, Record, WorkerDecisionSequenceLearn};

#[derive(Subcommand)]
pub enum LoraAction {
    /// Set up LoRA training environment (venv, dependencies, llama.cpp)
    Setup,
    /// List available adapters and GGUF files
    List,
    /// Generate sample training data for testing
    GenerateData {
        /// Number of samples to generate (default: 351)
        #[arg(short = 'n', long, default_value = "351")]
        samples: u32,
    },
    /// Prepare training data from SwarmEngine events
    PrepareData {
        /// Path to events.jsonl file
        #[arg(short, long)]
        events: Option<PathBuf>,

        /// Output JSONL file (default: lora/data/train.jsonl)
        #[arg(short, long)]
        output: Option<PathBuf>,

        /// Available actions (default: troubleshooting actions)
        #[arg(short, long, num_args = 1..)]
        actions: Option<Vec<String>>,
    },
    /// Train LoRA adapter from training data
    Train {
        /// Path to training data JSONL file (default: lora/data/train.jsonl)
        #[arg(short, long)]
        data: Option<PathBuf>,

        /// Output directory for adapter (default: lora/adapters/swarm-lora)
        #[arg(short, long)]
        output: Option<PathBuf>,

        /// Number of epochs (default: 3)
        #[arg(short, long, default_value = "3")]
        epochs: u32,

        /// LoRA rank (default: 16)
        #[arg(short, long, default_value = "16")]
        rank: u32,
    },
    /// Convert PEFT adapter to GGUF format
    Convert {
        /// Path to PEFT adapter directory (default: lora/adapters/swarm-lora)
        #[arg(short, long)]
        adapter: Option<PathBuf>,

        /// Output GGUF file (default: lora/gguf/<adapter-name>.gguf)
        #[arg(short, long)]
        output: Option<PathBuf>,
    },
    /// Show LoRA environment status
    Status,
    /// Clean generated files (adapters, gguf, cache)
    Clean {
        /// Also remove virtual environment
        #[arg(long)]
        all: bool,
    },
}

pub fn cmd_lora(action: LoraAction) {
    match action {
        LoraAction::Setup => cmd_lora_setup(),
        LoraAction::List => cmd_lora_list(),
        LoraAction::GenerateData { samples } => cmd_lora_generate_data(samples),
        LoraAction::PrepareData {
            events,
            output,
            actions,
        } => cmd_lora_prepare_data(events, output, actions),
        LoraAction::Train {
            data,
            output,
            epochs,
            rank,
        } => cmd_lora_train(data, output, epochs, rank),
        LoraAction::Convert { adapter, output } => cmd_lora_convert(adapter, output),
        LoraAction::Status => cmd_lora_status(),
        LoraAction::Clean { all } => cmd_lora_clean(all),
    }
}

/// Get path to lora directory (relative to project root or crate)
fn get_lora_dir() -> PathBuf {
    // Try to find lora/ directory relative to current directory
    let cwd = std::env::current_dir().expect("Failed to get current directory");

    // Check common locations
    let candidates = [
        cwd.join("lora"),
        cwd.join("..").join("..").join("lora"), // If running from target/debug
    ];

    for path in &candidates {
        if path.exists() {
            return path.canonicalize().unwrap_or_else(|_| path.clone());
        }
    }

    // Default: assume lora/ in current directory
    cwd.join("lora")
}

fn cmd_lora_setup() {
    let lora_dir = get_lora_dir();
    let setup_script = lora_dir.join("setup.sh");

    if !setup_script.exists() {
        eprintln!("Error: setup.sh not found at {}", setup_script.display());
        eprintln!("Make sure you're running from the project root.");
        std::process::exit(1);
    }

    println!("=== LoRA Setup ===");
    println!("Directory: {}", lora_dir.display());
    println!();

    let status = std::process::Command::new("bash")
        .arg(&setup_script)
        .current_dir(&lora_dir)
        .status()
        .expect("Failed to run setup.sh");

    if !status.success() {
        eprintln!("Setup failed with exit code: {:?}", status.code());
        std::process::exit(1);
    }
}

fn cmd_lora_train(data: Option<PathBuf>, output: Option<PathBuf>, epochs: u32, rank: u32) {
    let lora_dir = get_lora_dir();
    let venv_python = lora_dir.join(".venv").join("bin").join("python");
    let train_script = lora_dir.join("train.py");

    if !venv_python.exists() {
        eprintln!("Error: Virtual environment not found.");
        eprintln!("Run 'swarm-engine lora setup' first.");
        std::process::exit(1);
    }

    if !train_script.exists() {
        eprintln!("Error: train.py not found at {}", train_script.display());
        std::process::exit(1);
    }

    let mut cmd = std::process::Command::new(&venv_python);
    cmd.arg(&train_script)
        .arg("--epochs")
        .arg(epochs.to_string())
        .arg("--rank")
        .arg(rank.to_string())
        .current_dir(&lora_dir);

    if let Some(data_path) = data {
        cmd.arg("--data").arg(data_path);
    }
    if let Some(output_path) = output {
        cmd.arg("--output").arg(output_path);
    }

    println!("=== LoRA Training ===");
    println!("Epochs: {}, Rank: {}", epochs, rank);
    println!();

    let status = cmd.status().expect("Failed to run train.py");

    if !status.success() {
        eprintln!("Training failed with exit code: {:?}", status.code());
        std::process::exit(1);
    }
}

fn cmd_lora_convert(adapter: Option<PathBuf>, output: Option<PathBuf>) {
    let lora_dir = get_lora_dir();

    // Default paths
    let adapter_path = adapter.unwrap_or_else(|| lora_dir.join("adapters").join("swarm-lora"));
    let output_path = output.unwrap_or_else(|| {
        let adapter_name = adapter_path
            .file_name()
            .map(|n| n.to_string_lossy().to_string())
            .unwrap_or_else(|| "lora".to_string());
        lora_dir.join("gguf").join(format!("{}.gguf", adapter_name))
    });

    // Check adapter exists
    if !adapter_path.exists() {
        eprintln!("Error: Adapter not found: {}", adapter_path.display());
        eprintln!("Run 'swarm-engine lora train' first to create an adapter.");
        std::process::exit(1);
    }

    // Find llama.cpp convert script
    let convert_script = lora_dir.join("llama.cpp").join("convert_lora_to_gguf.py");
    if !convert_script.exists() {
        eprintln!(
            "Error: Conversion script not found: {}",
            convert_script.display()
        );
        eprintln!("Run 'swarm-engine lora setup' to clone llama.cpp.");
        std::process::exit(1);
    }

    // Create output directory
    if let Some(parent) = output_path.parent() {
        std::fs::create_dir_all(parent).ok();
    }

    // Default base model
    let base_model = "LiquidAI/LFM2.5-1.2B-Instruct";
    let dtype = "f16";

    println!("=== Converting LoRA to GGUF ===");
    println!("  Adapter: {}", adapter_path.display());
    println!("  Base: {}", base_model);
    println!("  Output: {}", output_path.display());
    println!("  Type: {}", dtype);
    println!();

    // Find Python - try python3 first, then python
    let python = find_python();

    let mut cmd = std::process::Command::new(&python);
    cmd.arg(&convert_script)
        .arg("--base-model-id")
        .arg(base_model)
        .arg("--outtype")
        .arg(dtype)
        .arg("--outfile")
        .arg(&output_path)
        .arg(&adapter_path)
        .current_dir(&lora_dir);

    println!(
        "Command: {} {} --base-model-id {} --outtype {} --outfile {} {}",
        python,
        convert_script.display(),
        base_model,
        dtype,
        output_path.display(),
        adapter_path.display()
    );
    println!();

    let status = cmd.status().expect("Failed to run convert_lora_to_gguf.py");

    if !status.success() {
        eprintln!(
            "\nError: Conversion failed with exit code: {:?}",
            status.code()
        );
        std::process::exit(1);
    }

    println!();
    println!("=== Conversion Complete ===");
    println!("GGUF adapter: {}", output_path.display());
    println!();
    println!("Usage with llama-server:");
    println!(
        "  llama-server -m <base-model.gguf> --lora {}",
        output_path.display()
    );
    println!();
    println!("Usage with swarm-engine:");
    println!("  cargo run --package swarm-engine-ui -- llama start \\");
    println!("    -m ~/.cache/.../LFM2.5-1.2B-Instruct-Q4_K_M.gguf \\");
    println!("    --lora {}", output_path.display());
}

/// Find Python executable (python3 or python)
fn find_python() -> String {
    // Try python3 first
    if std::process::Command::new("python3")
        .arg("--version")
        .output()
        .is_ok()
    {
        return "python3".to_string();
    }

    // Fall back to python
    if std::process::Command::new("python")
        .arg("--version")
        .output()
        .is_ok()
    {
        return "python".to_string();
    }

    eprintln!("Error: Python not found. Please install Python 3.");
    std::process::exit(1);
}

fn cmd_lora_status() {
    let lora_dir = get_lora_dir();

    println!("=== LoRA Environment Status ===");
    println!();

    // Check lora directory
    println!("Directory: {}", lora_dir.display());
    if !lora_dir.exists() {
        println!("  Status: NOT FOUND");
        println!();
        println!("Run 'swarm-engine lora setup' to initialize.");
        return;
    }
    println!("  Status: OK");
    println!();

    // Check virtual environment
    let venv_dir = lora_dir.join(".venv");
    print!("Virtual environment: ");
    if venv_dir.exists() {
        println!("OK ({})", venv_dir.display());
    } else {
        println!("NOT FOUND");
    }

    // Check llama.cpp
    let llama_cpp = lora_dir.join("llama.cpp");
    print!("llama.cpp: ");
    if llama_cpp.exists() {
        println!("OK");
    } else {
        println!("NOT FOUND");
    }

    // Check training data
    let train_data = lora_dir.join("data").join("train.jsonl");
    print!("Training data: ");
    if train_data.exists() {
        // Count lines
        if let Ok(content) = std::fs::read_to_string(&train_data) {
            let lines = content.lines().count();
            println!("{} samples", lines);
        } else {
            println!("OK");
        }
    } else {
        println!("NOT FOUND");
    }

    // Check adapters
    let adapters_dir = lora_dir.join("adapters");
    print!("Adapters: ");
    if adapters_dir.exists() {
        if let Ok(entries) = std::fs::read_dir(&adapters_dir) {
            let count = entries.filter(|e| e.is_ok()).count();
            if count > 0 {
                println!("{} adapter(s)", count);
            } else {
                println!("none");
            }
        } else {
            println!("OK");
        }
    } else {
        println!("none");
    }

    // Check GGUF files
    let gguf_dir = lora_dir.join("gguf");
    print!("GGUF files: ");
    if gguf_dir.exists() {
        if let Ok(entries) = std::fs::read_dir(&gguf_dir) {
            let gguf_files: Vec<_> = entries
                .filter_map(|e| e.ok())
                .filter(|e| e.path().extension().map(|x| x == "gguf").unwrap_or(false))
                .collect();
            if !gguf_files.is_empty() {
                println!("{} file(s)", gguf_files.len());
                for file in gguf_files {
                    println!("  - {}", file.file_name().to_string_lossy());
                }
            } else {
                println!("none");
            }
        } else {
            println!("OK");
        }
    } else {
        println!("none");
    }

    println!();
    if !venv_dir.exists() {
        println!("Run 'swarm-engine lora setup' to initialize environment.");
    }
}

fn cmd_lora_list() {
    let lora_dir = get_lora_dir();

    println!("=== LoRA Artifacts ===");
    println!();

    // List adapters
    let adapters_dir = lora_dir.join("adapters");
    println!("Adapters (PEFT format):");
    if adapters_dir.exists() {
        if let Ok(entries) = std::fs::read_dir(&adapters_dir) {
            let adapters: Vec<_> = entries
                .filter_map(|e| e.ok())
                .filter(|e| e.path().is_dir())
                .collect();
            if adapters.is_empty() {
                println!("  (none)");
            } else {
                for adapter in adapters {
                    let path = adapter.path();
                    let name = adapter.file_name();
                    // Check if it has adapter_model.safetensors
                    let has_model = path.join("adapter_model.safetensors").exists();
                    let status = if has_model { "ready" } else { "incomplete" };
                    println!("  - {} [{}]", name.to_string_lossy(), status);
                }
            }
        }
    } else {
        println!("  (none)");
    }

    println!();

    // List GGUF files
    let gguf_dir = lora_dir.join("gguf");
    println!("GGUF files (llama.cpp format):");
    if gguf_dir.exists() {
        if let Ok(entries) = std::fs::read_dir(&gguf_dir) {
            let gguf_files: Vec<_> = entries
                .filter_map(|e| e.ok())
                .filter(|e| e.path().extension().map(|x| x == "gguf").unwrap_or(false))
                .collect();
            if gguf_files.is_empty() {
                println!("  (none)");
            } else {
                for file in gguf_files {
                    let path = file.path();
                    let size = std::fs::metadata(&path)
                        .map(|m| format!("{:.1} MB", m.len() as f64 / 1_000_000.0))
                        .unwrap_or_else(|_| "?".to_string());
                    println!("  - {} ({})", file.file_name().to_string_lossy(), size);
                }
            }
        }
    } else {
        println!("  (none)");
    }

    println!();
    println!("Usage:");
    println!("  llama-server -m <base.gguf> --lora lora/gguf/<adapter>.gguf");
}

fn cmd_lora_generate_data(samples: u32) {
    let lora_dir = get_lora_dir();
    let venv_python = lora_dir.join(".venv").join("bin").join("python");
    let script = lora_dir.join("data").join("generate_training_data.py");

    if !venv_python.exists() {
        eprintln!("Error: Virtual environment not found.");
        eprintln!("Run 'swarm-engine lora setup' first.");
        std::process::exit(1);
    }

    if !script.exists() {
        eprintln!("Error: generate_training_data.py not found.");
        std::process::exit(1);
    }

    println!("=== Generating Training Data ===");
    println!("Samples: {}", samples);
    println!();

    // Note: The script generates a fixed number based on its logic
    // The samples parameter could be used to modify the script behavior
    let status = std::process::Command::new(&venv_python)
        .arg(&script)
        .current_dir(&lora_dir)
        .status()
        .expect("Failed to run generate_training_data.py");

    if !status.success() {
        eprintln!("Data generation failed with exit code: {:?}", status.code());
        std::process::exit(1);
    }
}

fn cmd_lora_prepare_data(
    events: Option<PathBuf>,
    output: Option<PathBuf>,
    actions: Option<Vec<String>>,
) {
    let lora_dir = get_lora_dir();

    // Default paths
    let events_dir = events.unwrap_or_else(|| lora_dir.join("data").join("raw"));
    let output_path = output.unwrap_or_else(|| lora_dir.join("data").join("train.jsonl"));

    println!("=== Preparing Training Data ===");
    println!("LearnModel: WorkerDecisionSequenceLearn");
    println!("Input: {}", events_dir.display());
    println!("Output: {}", output_path.display());
    println!();

    // Available actions (from --actions or defaults)
    let available_actions: Vec<String> = match actions {
        Some(a) => a,
        None => vec![
            "CheckStatus".to_string(),
            "ReadLogs".to_string(),
            "AnalyzeMetrics".to_string(),
            "Diagnose".to_string(),
            "Restart".to_string(),
        ],
    };

    // Create LearnModel
    let learn_model = WorkerDecisionSequenceLearn::new()
        .with_available_actions(available_actions.clone())
        .with_min_actions(3);

    // Find all .jsonl files
    let jsonl_files: Vec<PathBuf> = if events_dir.is_file() {
        vec![events_dir.clone()]
    } else if events_dir.is_dir() {
        find_jsonl_files(&events_dir)
    } else {
        eprintln!("Error: Input path not found: {}", events_dir.display());
        eprintln!("Run 'swarm-engine eval' first to generate events.");
        std::process::exit(1);
    };

    if jsonl_files.is_empty() {
        eprintln!("No .jsonl files found in {}", events_dir.display());
        std::process::exit(1);
    }

    // Process all files using LearnModel
    let mut all_records: Vec<Record> = Vec::new();
    let mut total_events = 0usize;

    for jsonl_path in &jsonl_files {
        println!(
            "Processing: {}",
            jsonl_path.file_name().unwrap_or_default().to_string_lossy()
        );

        let records = load_jsonl_as_records(jsonl_path);
        println!("  Records: {}", records.len());
        total_events += records.len();
        all_records.extend(records);
    }

    // Build Episodes using LearnModel
    let episodes = learn_model.build_episodes(&all_records);
    println!();
    println!("Episodes built: {}", episodes.len());

    // Convert to TrainingData
    let training_data: Vec<_> = episodes
        .iter()
        .filter_map(|ep| learn_model.convert(ep).ok())
        .collect();

    // Convert to Conversation format and deduplicate
    let mut seen = std::collections::HashSet::new();
    let unique_conversations: Vec<_> = training_data
        .iter()
        .map(|td| td.to_conversation())
        .filter(|conv| {
            let key = serde_json::to_string(conv).unwrap_or_default();
            seen.insert(key)
        })
        .collect();

    // Save output
    if let Some(parent) = output_path.parent() {
        std::fs::create_dir_all(parent).ok();
    }

    let mut file = std::fs::File::create(&output_path).expect("Failed to create output file");
    for conv in &unique_conversations {
        let json = serde_json::to_string(conv).expect("Failed to serialize");
        writeln!(file, "{}", json).expect("Failed to write");
    }

    println!();
    println!("=== Results ===");
    println!("Total events: {}", total_events);
    println!("Episodes: {}", episodes.len());
    println!("Training samples: {}", training_data.len());
    println!("Unique samples: {}", unique_conversations.len());
    println!("Saved to: {}", output_path.display());
}

// ============================================================================
// LoRA Data Preparation Helpers
// ============================================================================

/// Recursively find all .jsonl files in a directory
fn find_jsonl_files(dir: &PathBuf) -> Vec<PathBuf> {
    let mut files = Vec::new();

    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.filter_map(|e| e.ok()) {
            let path = entry.path();
            if path.is_dir() {
                files.extend(find_jsonl_files(&path));
            } else if path.extension().is_some_and(|ext| ext == "jsonl") {
                files.push(path);
            }
        }
    }

    files
}

/// Load JSONL file as Records (for LearnModel)
fn load_jsonl_as_records(path: &PathBuf) -> Vec<Record> {
    let file = match std::fs::File::open(path) {
        Ok(f) => f,
        Err(_) => return Vec::new(),
    };
    let reader = std::io::BufReader::new(file);

    reader
        .lines()
        .map_while(Result::ok)
        .filter(|line| !line.trim().is_empty())
        .enumerate()
        .filter_map(|(idx, line)| {
            let v: serde_json::Value = serde_json::from_str(&line).ok()?;

            let action = v.get("action")?.as_str()?.to_string();
            let target = v
                .get("target")
                .and_then(|t| t.as_str())
                .map(|s| s.to_string());
            let success = v.get("success")?.as_bool()?;

            // Create ActionRecord
            let mut record = ActionRecord::new(idx as u64, 0, action);
            if let Some(t) = target {
                record = record.target(t);
            }
            record = record.success(success);

            Some(Record::from(record))
        })
        .collect()
}

fn cmd_lora_clean(all: bool) {
    let lora_dir = get_lora_dir();

    println!("=== Cleaning LoRA Artifacts ===");
    println!();

    // Clean adapters
    let adapters_dir = lora_dir.join("adapters");
    if adapters_dir.exists() {
        print!("Removing adapters... ");
        if std::fs::remove_dir_all(&adapters_dir).is_ok() {
            println!("OK");
        } else {
            println!("FAILED");
        }
    }

    // Clean GGUF files
    let gguf_dir = lora_dir.join("gguf");
    if gguf_dir.exists() {
        print!("Removing GGUF files... ");
        if std::fs::remove_dir_all(&gguf_dir).is_ok() {
            println!("OK");
        } else {
            println!("FAILED");
        }
    }

    // Clean training data
    let train_data = lora_dir.join("data").join("train.jsonl");
    if train_data.exists() {
        print!("Removing training data... ");
        if std::fs::remove_file(&train_data).is_ok() {
            println!("OK");
        } else {
            println!("FAILED");
        }
    }

    // Optionally clean venv and llama.cpp
    if all {
        let venv_dir = lora_dir.join(".venv");
        if venv_dir.exists() {
            print!("Removing virtual environment... ");
            if std::fs::remove_dir_all(&venv_dir).is_ok() {
                println!("OK");
            } else {
                println!("FAILED");
            }
        }

        let llama_cpp = lora_dir.join("llama.cpp");
        if llama_cpp.exists() {
            print!("Removing llama.cpp... ");
            if std::fs::remove_dir_all(&llama_cpp).is_ok() {
                println!("OK");
            } else {
                println!("FAILED");
            }
        }
    }

    println!();
    println!("Clean complete.");
}