collet 0.1.1

use super::loader::load_config_file;
use super::paths::config_dir;
use super::types::{AgentDef, AgentsSection};
use super::wizard::RoleModels;
use std::path::Path;

// ---------------------------------------------------------------------------
// Agent discovery from .md files
// ---------------------------------------------------------------------------

/// Discover agent definitions. Priority (first match wins per agent name):
///   1. `<working_dir>/AGENTS.md`           — project single-file config
///   2. `<working_dir>/.claude/agents/*.md`  — project per-agent files
///   3. `~/.config/collet/AGENTS.md`         — user single-file config
///   4. `~/.config/collet/agents/*.md`        — user per-agent files
///
/// AGENTS.md format:
///   ## agent-name  model: default
///   Optional system prompt body.
pub fn discover_agents(working_dir: &str) -> Vec<AgentDef> {
    // Collect all user-defined agents from every source into a flat list.
    // Priority (first occurrence wins within this list):
    //   0. config.toml [[agents.list]]
    //   1. Project AGENTS.md
    //   2. Project .claude/agents/
    //   3. User ~/.collet/AGENTS.md
    //   4. User ~/.collet/agents/
    let mut user_agents: Vec<AgentDef> = Vec::new();
    let mut seen = std::collections::HashSet::new();

    // 0. config.toml — [[agents.list]] entries first, then [agents] code/architect/ask
    //    shorthand fields.  The shorthand fields are written by `collet setup` and
    //    must feed into discovery so the correct per-role model is used at runtime.
    if let Ok(config_file) = load_config_file() {
        for agent_def in config_file.agents.list {
            if seen.insert(agent_def.name.clone()) {
                user_agents.push(agent_def);
            }
        }

        // [agents] code/architect/ask shorthand → synthesize AgentDef entries so
        // that the role-specific models are respected even without explicit .md files.
        //
        // Use `providers` (Vec<String>) rather than `provider` (Option<String>) because
        // switch_agent() resolves the active model from `agent.providers` via find_map.
        // A bare `provider` field is ignored there, which caused all three roles to fall
        // back to resolve_default_provider() and show the global default model.
        for (role, val) in [
            ("arbor", &config_file.agents.arbor),
            ("architect", &config_file.agents.architect),
            ("code", &config_file.agents.code),
            ("ask", &config_file.agents.ask),
        ] {
            if seen.contains(role) {
                continue;
            }
            let Some(prov_model) = val else { continue };
            // prov_model is already in "provider/model" format from `collet setup`.
            // Store it verbatim in `providers` so switch_agent() can resolve it.
            let (provider, model) = if let Some(slash) = prov_model.find('/') {
                (
                    Some(prov_model[..slash].to_string()),
                    prov_model[slash + 1..].to_string(),
                )
            } else {
                (None, prov_model.clone())
            };
            seen.insert(role.to_string());
            user_agents.push(AgentDef {
                name: role.to_string(),
                model,
                provider,
                providers: vec![prov_model.clone()],
                ..Default::default()
            });
        }
    }

    // 1. Project AGENTS.md (highest priority)
    load_agents_from_md_file(
        &Path::new(working_dir).join("AGENTS.md"),
        &mut user_agents,
        &mut seen,
    );

    // 2. Project per-agent files
    load_agents_from_dir(
        &Path::new(working_dir).join(".claude").join("agents"),
        &mut user_agents,
        &mut seen,
    );

    // 3. User AGENTS.md
    load_agents_from_md_file(&config_dir().join("AGENTS.md"), &mut user_agents, &mut seen);

    // 4. User per-agent files
    load_agents_from_dir(&config_dir().join("agents"), &mut user_agents, &mut seen);

    // ── Assemble final list ──────────────────────────────────────────────────
    // Built-in core roles (arbor → architect → code → ask) are ALWAYS first
    // so they are immediately visible at the top of the agent picker.
    //
    // If the user defined an agent with the same name, their version replaces
    // the built-in default (customisation is preserved).  Any remaining
    // user-defined agents follow in discovery order.
    let mut result: Vec<AgentDef> = Vec::new();
    for builtin in AgentsSection::default().list {
        if let Some(pos) = user_agents.iter().position(|a| a.name == builtin.name) {
            let mut agent = user_agents.remove(pos);
            // Inherit tier from builtin when the user's version has none.
            // [agents] shorthand entries (ask/code/architect/arbor) are synthesized
            // without tier, so without this the router ignores them and falls back
            // to the first domain agent that declares `tier: light/heavy`.
            if agent.tier.is_none() {
                agent.tier = builtin.tier;
            }
            result.push(agent);
        } else {
            result.push(builtin);
        }
    }
    result.extend(user_agents);

    tracing::info!(count = result.len(), "Agent definitions loaded");
    result
}

/// Parse a single AGENTS.md containing agent role definitions.
/// Supports both `## name  model: <model>` and `### name  model: <model>` formats.
/// Non-agent sections (Commands, Boundaries, etc.) are ignored here —
/// they are loaded separately by `prompt::load_agents_md_context()`.
fn load_agents_from_md_file(
    path: &Path,
    agents: &mut Vec<AgentDef>,
    seen: &mut std::collections::HashSet<String>,
) {
    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(_) => return,
    };
    for line in content.lines() {
        // Support both ## and ### headings with "model:" marker
        let rest = line
            .trim()
            .strip_prefix("### ")
            .or_else(|| line.trim().strip_prefix("## "));
        let Some(rest) = rest else { continue };
        // Only parse lines containing "model:" — skip non-agent sections
        let Some(idx) = rest.find("model:") else {
            continue;
        };
        let name = rest[..idx].trim().to_string();
        let model = rest[idx + 6..].trim().to_string();
        if !name.is_empty() && !model.is_empty() && seen.insert(name.clone()) {
            agents.push(AgentDef {
                name,
                model,
                provider: None,
                system_prompt: String::new(),
                ..AgentDef::default()
            });
        }
    }
}

fn load_agents_from_dir(
    dir: &Path,
    agents: &mut Vec<AgentDef>,
    seen: &mut std::collections::HashSet<String>,
) {
    let entries = match std::fs::read_dir(dir) {
        Ok(e) => e,
        Err(_) => return,
    };

    for entry in entries.flatten() {
        let path = entry.path();
        if path.extension().and_then(|e| e.to_str()) != Some("md") {
            continue;
        }
        if let Some(agent) = parse_agent_md(&path)
            && seen.insert(agent.name.clone())
        {
            agents.push(agent);
        }
    }
}

pub(crate) fn parse_agent_md(path: &Path) -> Option<AgentDef> {
    let content = std::fs::read_to_string(path).ok()?;

    // Extract YAML frontmatter
    let trimmed = content.trim_start();
    if !trimmed.starts_with("---") {
        return None;
    }
    let after_first = &trimmed[3..];
    let end = after_first.find("---")?;
    let frontmatter = &after_first[..end];

    let name = extract_yaml_value(frontmatter, "name")?;
    let description = extract_yaml_value(frontmatter, "description");
    let tier = extract_yaml_value(frontmatter, "tier").and_then(|v| match v.as_str() {
        "heavy" | "medium" | "light" => Some(v),
        other => {
            tracing::warn!(agent = %name, tier = %other, "Unknown tier value, ignoring");
            None
        }
    });
    let raw_model = extract_yaml_value(frontmatter, "model").unwrap_or_default();

    // Support "provider/model" shorthand in the model field.
    // If the value contains exactly one '/' and no explicit `provider:` is set,
    // split it so callers get the correct provider context automatically.
    let explicit_provider = extract_yaml_value(frontmatter, "provider");
    let (model, provider) = if explicit_provider.is_none() {
        if let Some(slash) = raw_model.find('/') {
            let p = raw_model[..slash].to_string();
            let m = raw_model[slash + 1..].to_string();
            (m, Some(p))
        } else {
            (raw_model, None)
        }
    } else {
        (raw_model, explicit_provider)
    };
    let providers: Vec<String> = extract_yaml_value(frontmatter, "providers")
        .map(|v| {
            v.split(',')
                .map(|s| s.trim().to_string())
                .filter(|s| !s.is_empty())
                .collect()
        })
        .unwrap_or_default();

    // When `model:` and `provider:` are absent (providers-only frontmatter), derive
    // them from the first entry in `providers` (format: "provider/model").
    let (model, provider) = if model.is_empty() && provider.is_none() {
        if let Some(first) = providers.first() {
            if let Some(slash) = first.find('/') {
                (
                    first[slash + 1..].to_string(),
                    Some(first[..slash].to_string()),
                )
            } else {
                (model, provider)
            }
        } else {
            (model, provider)
        }
    } else {
        (model, provider)
    };

    // Support tags: or categories: (alias), all three YAML formats
    let tags: Vec<String> = {
        let t = extract_yaml_list(frontmatter, "tags");
        if t.is_empty() {
            extract_yaml_list(frontmatter, "categories")
        } else {
            t
        }
    };
    let temperature: Option<f32> =
        extract_yaml_value(frontmatter, "temperature").and_then(|v| v.parse().ok());
    let thinking_budget_tokens: Option<u32> =
        extract_yaml_value(frontmatter, "thinking_budget_tokens").and_then(|v| v.parse().ok());
    let reasoning_effort = extract_yaml_value(frontmatter, "reasoning_effort");
    let max_iterations: Option<u32> =
        extract_yaml_value(frontmatter, "max_iterations").and_then(|v| v.parse().ok());
    let iteration_delay_ms: Option<u64> =
        extract_yaml_value(frontmatter, "iteration_delay_ms").and_then(|v| v.parse().ok());
    let max_output_tokens: Option<u32> =
        extract_yaml_value(frontmatter, "max_output_tokens").and_then(|v| v.parse().ok());
    let supports_tools: Option<bool> =
        extract_yaml_value(frontmatter, "supports_tools").and_then(|v| v.parse().ok());
    let supports_reasoning: Option<bool> =
        extract_yaml_value(frontmatter, "supports_reasoning").and_then(|v| v.parse().ok());
    let context_window: Option<usize> =
        extract_yaml_value(frontmatter, "context_window").and_then(|v| v.parse().ok());
    let soul: Option<bool> = extract_yaml_value(frontmatter, "soul").and_then(|v| v.parse().ok());

    // Body after closing `---` is the agent system prompt
    let body = after_first[end + 3..].trim().to_string();

    Some(AgentDef {
        name,
        description,
        tier,
        model,
        provider,
        providers,
        tags,
        system_prompt: body,
        temperature,
        thinking_budget_tokens,
        reasoning_effort,
        max_iterations,
        iteration_delay_ms,
        max_output_tokens,
        supports_tools,
        supports_reasoning,
        context_window,
        soul,
    })
}

fn extract_yaml_value(yaml: &str, field: &str) -> Option<String> {
    for line in yaml.lines() {
        let line = line.trim();
        if let Some(rest) = line.strip_prefix(field) {
            let rest = rest.trim_start();
            if let Some(value) = rest.strip_prefix(':') {
                let value = value.trim().trim_matches('"').trim_matches('\'');
                if !value.is_empty() {
                    return Some(value.to_string());
                }
            }
        }
    }
    None
}

/// Parse a YAML frontmatter list field that may be written in three formats:
/// - Comma-separated string:  `tags: foo, bar, baz`
/// - Inline YAML array:       `tags: [foo, bar, baz]`
/// - Block YAML array:        `tags:\n  - foo\n  - bar`
///
/// Also accepts `categories` as an alias for `tags`.
pub fn extract_yaml_list(yaml: &str, field: &str) -> Vec<String> {
    let lines: Vec<&str> = yaml.lines().collect();
    for (i, line) in lines.iter().enumerate() {
        let trimmed = line.trim();
        let Some(rest) = trimmed.strip_prefix(field) else {
            continue;
        };
        let rest = rest.trim_start();
        let Some(after_colon) = rest.strip_prefix(':') else {
            continue;
        };
        let value = after_colon.trim();

        if value.is_empty() {
            // Block array: collect subsequent `- item` lines
            return lines[i + 1..]
                .iter()
                .take_while(|l| {
                    let t = l.trim();
                    t.starts_with("- ") || t.is_empty()
                })
                .filter_map(|l| {
                    l.trim()
                        .strip_prefix("- ")
                        .map(|v| v.trim().trim_matches('"').trim_matches('\'').to_lowercase())
                })
                .filter(|s| !s.is_empty())
                .collect();
        } else if value.starts_with('[') && value.ends_with(']') {
            // Inline YAML array: [foo, bar, "baz qux"]
            let inner = &value[1..value.len() - 1];
            return inner
                .split(',')
                .map(|s| s.trim().trim_matches('"').trim_matches('\'').to_lowercase())
                .filter(|s| !s.is_empty())
                .collect();
        } else {
            // Comma-separated string
            return value
                .split(',')
                .map(|s| s.trim().trim_matches('"').trim_matches('\'').to_lowercase())
                .filter(|s| !s.is_empty())
                .collect();
        }
    }
    vec![]
}

/// Update model and provider fields in an agent .md file's YAML frontmatter.
/// Preserves all other frontmatter fields and the body content.
/// Apply role-based model assignments to ALL agent .md files in `agents_dir`.
///
/// Each agent's `mode:` frontmatter field determines which role model to use:
///   - "architect" → `roles.architect`
///   - "ask"       → `roles.ask`
///   - anything else (including "code", unset) → `roles.code`
///
/// This ensures domain agents (devops, security, finance, etc.) also receive
/// a real model instead of staying on the generic "default" alias.
pub fn apply_roles_to_agents(agents_dir: &Path, roles: &RoleModels, provider_name: &str) {
    let Ok(entries) = std::fs::read_dir(agents_dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.extension().and_then(|e| e.to_str()) != Some("md") {
            continue;
        }
        let agent = match parse_agent_md(&path) {
            Some(a) => a,
            None => continue,
        };
        // architect/code/ask are managed in config.toml — never touch their .md files.
        if matches!(agent.name.as_str(), "architect" | "code" | "ask") {
            continue;
        }
        // Only update agents that have no provider set yet (freshly scaffolded).
        // Existing hand-edited assignments are preserved.
        let agent_prov = agent.provider.as_deref().unwrap_or("");
        if !agent_prov.is_empty() {
            continue;
        }
        let model = &roles.code;
        update_agent_md_frontmatter(&path, model, provider_name);
    }
}

fn update_agent_md_frontmatter(path: &Path, model: &str, provider: &str) {
    update_agent_md_frontmatter_full(path, model, provider, None);
}

/// Writes `providers:` as the sole provider/model field in agent frontmatter.
/// Resolution order at runtime: `providers[0]` → `providers[1]` → … → global `[default]`.
///
/// `providers_chain`: comma-separated "provider/model" pairs for fallback chain.
/// When None, a single-entry chain is synthesised from `provider`/`model`.
pub(crate) fn update_agent_md_frontmatter_full(
    path: &Path,
    model: &str,
    provider: &str,
    providers_chain: Option<&str>,
) {
    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(_) => return,
    };
    let trimmed = content.trim_start();
    if !trimmed.starts_with("---") {
        return;
    }
    let after_first = &trimmed[3..];
    let end = match after_first.find("---") {
        Some(e) => e,
        None => return,
    };
    let frontmatter = &after_first[..end];
    // body_text = everything after the closing "---", leading newlines stripped.
    let body_text = after_first[end + 3..].trim_start_matches(['\n', '\r']);

    // `providers:` is the single source of truth for provider/model in agent files.
    // model/provider are encoded per-entry as "provider/model".
    // Runtime fallback: providers[0] → providers[1] → … → global [default].
    //
    // When no explicit chain is given, synthesise a single-entry chain from the params.
    let chain: String = match providers_chain {
        Some(c) if !c.is_empty() => c.to_string(),
        _ => format!("{provider}/{model}"),
    };

    // Rebuild frontmatter: drop model/provider fields, replace/insert providers.
    // Use trim() to avoid accumulating leading blank lines across repeated setup runs.
    let mut new_lines: Vec<String> = Vec::new();
    let mut has_providers = false;
    for line in frontmatter.trim().lines() {
        let trimmed_line = line.trim();
        if trimmed_line.starts_with("model:")
            || trimmed_line.starts_with("model :")
            || trimmed_line.starts_with("provider:")
            || trimmed_line.starts_with("provider :")
        {
            // drop — encoded in providers entries
        } else if trimmed_line.starts_with("providers:") || trimmed_line.starts_with("providers :")
        {
            new_lines.push(format!("providers: {chain}"));
            has_providers = true;
        } else {
            new_lines.push(line.to_string());
        }
    }
    if !has_providers {
        let pos = new_lines
            .iter()
            .position(|l| l.trim().starts_with("name:"))
            .map(|i| i + 1)
            .unwrap_or(0);
        new_lines.insert(pos, format!("providers: {chain}"));
    }

    let result = if body_text.is_empty() {
        format!("---\n{}\n---\n", new_lines.join("\n"))
    } else {
        format!("---\n{}\n---\n\n{}", new_lines.join("\n"), body_text)
    };
    let _ = std::fs::write(path, result);
}