planner / src/main.rs

raw file 15574 bytes
mod cli;
mod knowledge;
mod llm;
mod model;
mod planner;
mod prompt;
mod templates;

use anyhow::Result;
use clap::Parser;
use cli::{Cli, Commands};
use colored::Colorize;
use dialoguer::{theme::ColorfulTheme, Input, Select};
use llm::{Backend, LlmConfig};
use model::ProjectPlan;
use std::collections::BTreeMap;
use std::fs;
use std::path::Path;

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.command {
        None => {
            println!("{}", "Planner — interactive mode".bold());
            run_wizard(None, ".".to_string(), cli.llm_ip, cli.llm_port)?;
        }
        Some(Commands::New { name, out, llm_ip, llm_port }) => {
            run_wizard(name, out, llm_ip, llm_port)?;
        }
        Some(Commands::Templates) => {
            for t in templates::list_templates() {
                println!("{:<14} {} [{}]", t.manifest.slug, t.manifest.description, t.manifest.category);
            }
        }
        Some(Commands::Show { slug }) => {
            if let Some(t) = templates::find_template(&slug) {
                println!("{} — {}", t.manifest.name, t.manifest.description);
                for (p, _) in &t.files {
                    println!("  {}", p);
                }
            } else {
                eprintln!("template not found: {}", slug);
            }
        }
        Some(Commands::Init { template, name, out }) => {
            do_init(&template, name, &out)?;
        }
        Some(Commands::Plan { idea, r#type, language, out, llm_ip, llm_port, constraints }) => {
            let (db, auth, deploy, scale, mut features) = planner::infer_requirements(&idea);
            let ptype = r#type.clone().unwrap_or_else(|| "cli".to_string());
            let template_slug = planner::template_slug_for(&knowledge::project_type_aliases(&ptype)).map(|s| s.to_string());

            let cfg_opt = llm_ip.as_ref().map(|ip| LlmConfig::new(llm::Backend::Llamacpp, None, None, Some(format!("http://{}:{}", ip, llm_port))));

            if let Some(cfg) = &cfg_opt {
                // suggest features if empty or refine
                if features.is_empty() {
                    let prompt = format!("Suggest 8-12 concrete features for a {} project: {}. Return ONLY a JSON array of strings.", ptype, idea);
                    if let Ok(out) = llm::generate(cfg, &prompt) {
                        // try JSON
                        let mut parsed = Vec::new();
                        if let Some(s) = out.find('[') {
                            if let Some(e) = out.rfind(']') {
                                if e > s {
                                    if let Ok(arr) = serde_json::from_str::<Vec<String>>(&out[s..=e]) {
                                        parsed = arr.into_iter().map(|x| x.trim().to_string()).filter(|x| !x.is_empty()).collect();
                                    }
                                }
                            }
                        }
                        if parsed.is_empty() {
                            parsed = llm::parse_feature_lines(&out);
                        }
                        if !parsed.is_empty() { features = parsed; }
                    }
                } else {
                    if let Ok(refined) = llm::refine_features(cfg, &features, &ptype, &idea) {
                        features = refined;
                    }
                }
            }

            let input = planner::PlanInput {
                name: slugify(&idea.chars().take(30).collect::<String>()),
                description: idea.clone(),
                project_type_raw: ptype,
                audience: "developers".to_string(),
                scale,
                language_pref: language,
                database: db,
                auth,
                deployment: deploy,
                features,
                constraints: constraints.map(|s| s.split(',').map(|x| x.trim().to_string()).filter(|x| !x.is_empty()).collect()).unwrap_or_default(),
                template_slug,
            };
            let mut plan = planner::build_plan(input.clone());

            if let Some(cfg) = &cfg_opt {
                println!("{}", "Consulting local model...".dimmed());
                match llm::plan_assistance(cfg, &input, &plan) {
                    Ok(notes) => {
                        if notes.contains("## ") {
                            // merge risks from LLM into plan.risks
                            if let Some(start) = notes.find("## Risks to Watch") {
                                let header = "## Risks to Watch";
                                let rest = &notes[start + header.len()..];
                                let end = rest.find("## ").unwrap_or(rest.len());
                                let section = &rest[..end];
                                for line in section.lines() {
                                    let l = line.trim().trim_start_matches(['-', '*', ' ']).trim().to_string();
                                    if !l.is_empty() && !l.starts_with('#') && !plan.risks.contains(&l) {
                                        plan.risks.push(l);
                                    }
                                }
                            }
                            plan.ai_assistance = Some(notes);
                        } else {
                            eprintln!("local model returned no structured sections; skipping.");
                        }
                    }
                    Err(e) => eprintln!("local model failed: {}", e),
                }
            }
            write_plan(&plan, Path::new(&out))?;
        }
    }
    Ok(())
}

fn run_wizard(name_opt: Option<String>, out: String, llm_ip: Option<String>, llm_port: u16) -> Result<()> {
    let theme = ColorfulTheme::default();

    let mut backend = Backend::Auto;
    let mut model: Option<String> = None;
    let ollama_url: Option<String> = None;
    let mut llama_url: Option<String> = None;

    let use_local = if let Some(ip) = &llm_ip {
        backend = Backend::Llamacpp;
        llama_url = Some(format!("http://{}:{}", ip, llm_port));
        println!("{}", format!("Using LAN LLM at http://{}:{}", ip, llm_port).green());
        true
    } else {
        let use_local = Select::with_theme(&theme)
            .with_prompt("Use a local model to assist with planning?")
            .items(&["no", "yes"])
            .default(0)
            .interact()? == 1;

        if use_local {
            let backends = ["auto", "ollama", "llamacpp"];
            let bi = Select::with_theme(&theme)
                .with_prompt("Backend")
                .items(&backends)
                .default(0)
                .interact()?;
            backend = Backend::parse(backends[bi]);
            let m: String = Input::with_theme(&theme)
                .with_prompt("Model name (blank = default)")
                .default(String::new())
                .interact_text()?;
            if !m.trim().is_empty() {
                model = Some(m);
            }
        }
        use_local
    };

    let name: String = match name_opt {
        Some(n) => n,
        None => Input::with_theme(&theme)
            .with_prompt("Project name")
            .default("my-project".to_string())
            .interact_text()?,
    };

    let description: String = Input::with_theme(&theme)
        .with_prompt("One-line description / idea")
        .interact_text()?;

    let types = ["cli", "rest-api", "web-app", "library", "desktop", "mobile", "data-pipeline", "ml-project", "game", "other"];
    let ti = Select::with_theme(&theme)
        .with_prompt("Project type")
        .items(&types)
        .default(0)
        .interact()?;
    let project_type_raw = types[ti].to_string();

    let audience: String = Input::with_theme(&theme)
        .with_prompt("Target audience")
        .default("developers".to_string())
        .interact_text()?;

    let scales = ["prototype", "small", "medium", "large"];
    let si = Select::with_theme(&theme).with_prompt("Scale").items(&scales).default(1).interact()?;
    let scale = scales[si].to_string();

    let langs = ["auto", "rust", "go", "python", "typescript", "java", "kotlin", "csharp", "swift"];
    let li = Select::with_theme(&theme).with_prompt("Language preference").items(&langs).default(0).interact()?;
    let language_pref = if langs[li] == "auto" { None } else { Some(langs[li].to_string()) };

    let dbs = ["none", "sqlite", "postgres", "mongo"];
    let di = Select::with_theme(&theme).with_prompt("Database").items(&dbs).default(0).interact()?;
    let database = dbs[di].to_string();

    let auths = ["none", "basic", "oauth", "api-key"];
    let ai = Select::with_theme(&theme).with_prompt("Auth").items(&auths).default(0).interact()?;
    let auth = auths[ai].to_string();

    let deploys = ["local", "cloud", "container", "serverless", "desktop"];
    let dpi = Select::with_theme(&theme).with_prompt("Deployment target").items(&deploys).default(0).interact()?;
    let deployment = deploys[dpi].to_string();

    // Feature planning
    println!("{}", "Plan features (comma-separated). Examples will vary by project type.".dimmed());
    let features_str: String = Input::with_theme(&theme)
        .with_prompt("Features")
        .default(String::new())
        .interact_text()?;
    let mut features: Vec<String> = features_str.split(',').map(|s| s.trim().to_string()).filter(|s| !s.is_empty()).collect();

    if use_local && features.is_empty() {
        let cfg = LlmConfig::new(backend, model.clone(), ollama_url.clone(), llama_url.clone());
        let prompt = format!("Suggest 8-12 concrete features for a {} project: {}. Return one feature per line, no extra text.", project_type_raw, description);
        if let Ok(out) = llm::generate(&cfg, &prompt) {
            for line in out.lines() {
                let l = line.trim().trim_start_matches(['-', '*', '•', ' ']).trim().to_string();
                if !l.is_empty() { features.push(l); }
            }
            if !features.is_empty() {
                println!("{}", "Suggested features:".green());
                for f in &features { println!("  - {}", f); }
            }
        }
    }

    if use_local && !features.is_empty() {
        let cfg = LlmConfig::new(backend, model.clone(), ollama_url.clone(), llama_url.clone());
        println!("{}", "Refining features...".dimmed());
        if let Ok(refined) = llm::refine_features(&cfg, &features, &project_type_raw, &description) {
            println!("{}", "Refined features:".green());
            for f in &refined {
                println!("  - {}", f);
            }
            features = refined;
        }
    }

    let constraints_str: String = Input::with_theme(&theme)
        .with_prompt("Constraints (comma-separated, optional)")
        .default(String::new())
        .interact_text()?;
    let constraints: Vec<String> = constraints_str.split(',').map(|s| s.trim().to_string()).filter(|s| !s.is_empty()).collect();

    let templates_list = templates::list_templates();
    let mut tmpl_choices: Vec<String> = vec!["(none)".to_string()];
    for t in &templates_list {
        tmpl_choices.push(format!("{} — {}", t.manifest.slug, t.manifest.description));
    }
    let tsi = Select::with_theme(&theme)
        .with_prompt("Template to reference in plan (optional)")
        .items(&tmpl_choices)
        .default(0)
        .interact()?;
    let template_slug = if tsi == 0 { None } else { Some(templates_list[tsi - 1].manifest.slug.clone()) };

    let input = planner::PlanInput {
        name: name.clone(),
        description,
        project_type_raw,
        audience,
        scale,
        language_pref,
        database,
        auth,
        deployment,
        features: features.clone(),
        constraints,
        template_slug: template_slug.clone(),
    };
    let input_for_llm = input.clone();
    let mut plan = planner::build_plan(input);

    if use_local {
        let cfg = LlmConfig::new(backend, model, ollama_url, llama_url);
        println!("{}", "Consulting local model...".dimmed());
        // feature suggestions
        let feat_prompt = format!(
            "Suggest 8-12 concrete features for a {} project described as: {}. Return a bullet list only.",
            plan.project_type, plan.description
        );
        if let Ok(sugg) = llm::generate(&cfg, &feat_prompt) {
            // append suggested features as a note in ai_assistance later; store temporarily
            plan.ai_assistance = Some(format!("## Suggested Features\n{}\n", sugg));
        }
        match llm::plan_assistance(&cfg, &input_for_llm, &plan) {
            Ok(notes) => {
                if notes.contains("## ") {
                    if let Some(start) = notes.find("## Risks to Watch") {
                        let header = "## Risks to Watch";
                        let rest = &notes[start + header.len()..];
                        let end = rest.find("## ").unwrap_or(rest.len());
                        let section = &rest[..end];
                        for line in section.lines() {
                            let l = line.trim().trim_start_matches(['-', '*', ' ']).trim().to_string();
                            if !l.is_empty() && !l.starts_with('#') && !plan.risks.contains(&l) {
                                plan.risks.push(l);
                            }
                        }
                    }
                    let combined = match &plan.ai_assistance {
                        Some(prev) => format!("{}\n\n{}", prev, notes),
                        None => notes,
                    };
                    plan.ai_assistance = Some(combined);
                } else {
                    eprintln!("local model returned no structured sections; skipping.");
                }
            }
            Err(e) => eprintln!("local model failed: {}", e),
        }
    }

    let out_path = Path::new(&out);
    fs::create_dir_all(out_path)?;

    write_plan(&plan, out_path)?;
    println!("{}", "Done. PLAN.md written — feed it to your LLM.".green().bold());
    Ok(())
}

fn do_init(template: &str, name: Option<String>, out: &str) -> Result<()> {
    let t = templates::find_template(template).ok_or_else(|| anyhow::anyhow!("template not found: {}", template))?;
    let name = name.unwrap_or_else(|| template.to_string());
    let mut vars = BTreeMap::new();
    vars.insert("name".to_string(), slugify(&name));
    vars.insert("description".to_string(), format!("{} project", name));
    let out_path = Path::new(out);
    fs::create_dir_all(out_path)?;
    let written = templates::scaffold_template(&t, &vars, out_path)?;
    for w in written {
        println!("{}", w);
    }
    Ok(())
}

fn write_plan(plan: &ProjectPlan, out: &Path) -> Result<()> {
    fs::create_dir_all(out)?;
    let content = prompt::render_prompt(plan);
    fs::write(out.join("PLAN.md"), content)?;
    let json = serde_json::to_string_pretty(plan)?;
    fs::write(out.join("plan.json"), json)?;
    println!("{}", format!("Wrote {}/PLAN.md", out.display()));
    Ok(())
}

fn slugify(s: &str) -> String {
    let mut out = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            out.push(c.to_ascii_lowercase());
        } else if c == ' ' || c == '-' || c == '_' {
            if !out.ends_with('-') {
                out.push('-');
            }
        }
    }
    let out = out.trim_matches('-').to_string();
    if out.is_empty() { "project".to_string() } else { out }
}