procyon 0.1.2

use std::sync::Arc;
use std::time::Duration;

use crate::agent::subagent::{self, SubAgentConfig};
use crate::config::AppConfig;
use crate::personas::Persona;

const MAX_CONCURRENT_SUBAGENTS: usize = 5;

/// The share of a persona's whole budget that any single request may take.
///
/// Four fifths leaves room for the outer timeout to still be the one that catches a persona which
/// is slow across many rounds, while a single stalled request trips the inner one first and gets
/// to say which round hung.
fn request_share(total: Duration) -> Duration {
    (total / 5 * 4).max(Duration::from_secs(1))
}

/// Response from a single persona in party mode.
pub struct PartyResponse {
    pub persona: Persona,
    pub text: String,
    pub round_trips: usize,
    pub elapsed: Duration,
    pub error: Option<String>,
}

/// Configuration for a party mode discussion.
pub struct PartyConfig {
    /// The message to discuss.
    pub message: String,
    /// Specific personas to include (None = all available).
    pub personas: Option<Vec<String>>,
    /// Model override for all sub-agents.
    pub model: Option<String>,
    /// Timeout per sub-agent (None = 120 seconds).
    pub timeout: Option<Duration>,
}

/// Everything one persona's run shares with the others in the same party.
///
/// Grouped rather than passed one by one: only the persona itself varies per run, so a parameter
/// list made every call site restate seven values that are the same for all of them.
struct PartyRun<'a> {
    config: &'a AppConfig,
    message: &'a str,
    model: &'a Option<String>,
    workspace_context: &'a str,
    timeout: Duration,
    semaphore: Arc<tokio::sync::Semaphore>,
}

/// Runs a single persona sub-agent and returns its response.
async fn run_persona(run: &PartyRun<'_>, persona: &Persona) -> PartyResponse {
    let timeout = run.timeout;
    let _permit = run.semaphore.acquire().await.expect("semaphore closed");
    let start = std::time::Instant::now();
    let system_prompt = persona.system_prompt(run.workspace_context);

    // Each seat at the table gets its own registry, cut to its own tools and ceiling — the point
    // of running distinct personas is that they differ operationally, not only in voice. A shared
    // registry gave every seat the same capabilities regardless of who was sitting in it.
    let loaded = crate::runtime::persona_tools(run.config, persona).await;

    let subagent_config = SubAgentConfig {
        system_prompt,
        message: run.message.to_string(),
        model: run.model.clone(),
        max_tokens: None,
        max_rounds: None,
        allowed_tools: None,
        // Strictly smaller than the outer timeout, which is the point: the two budgets measure
        // different things — one request against the whole persona — and setting them equal meant
        // the outer one always won. The per-round message never appeared, and the text produced
        // before the stall was discarded with it.
        timeout_secs: Some(request_share(timeout).as_secs()),
    };

    let result = tokio::time::timeout(
        timeout,
        subagent::run_subagent(run.config, subagent_config, &loaded.registry),
    )
    .await;

    let elapsed = start.elapsed();
    let persona = persona.clone();
    match result {
        Ok(Ok(response)) => PartyResponse {
            persona,
            text: response.text,
            round_trips: response.round_trips,
            elapsed,
            error: None,
        },
        Ok(Err(e)) => PartyResponse {
            persona,
            text: String::new(),
            round_trips: 0,
            elapsed,
            error: Some(format!("Failed: {}", e)),
        },
        Err(_) => PartyResponse {
            persona,
            text: String::new(),
            round_trips: 0,
            elapsed,
            error: Some(format!("Timed out after {:?}", timeout)),
        },
    }
}

/// Runs a party mode discussion with multiple personas in parallel.
///
/// Each persona runs as an independent sub-agent.
/// Responses are collected with a per-persona timeout.
/// Concurrency is limited to MAX_CONCURRENT_SUBAGENTS.
pub async fn run_party(config: &AppConfig, party_config: PartyConfig) -> Vec<PartyResponse> {
    let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));

    // Shared and already discovered off the runtime: party mode used to walk every skill
    // directory on entry, on a worker thread, once per invocation.
    let persona_reg = crate::registries::personas().await;

    // Select personas
    let selected: Vec<Persona> = match &party_config.personas {
        Some(names) => names
            .iter()
            .filter_map(|name| persona_reg.get(name).cloned())
            .collect(),
        None => persona_reg.all().to_vec(),
    };

    if selected.is_empty() {
        return Vec::new();
    }

    let timeout = party_config.timeout.unwrap_or(Duration::from_secs(120));

    let workspace_context = format!("Workspace: {}", cwd.display());
    let run = PartyRun {
        config,
        message: &party_config.message,
        model: &party_config.model,
        workspace_context: &workspace_context,
        timeout,
        semaphore: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_SUBAGENTS)),
    };

    // Run all personas concurrently with bounded parallelism
    let futures: Vec<_> = selected
        .iter()
        .map(|persona| run_persona(&run, persona))
        .collect();

    futures_util::future::join_all(futures).await
}

/// Formats party responses into a readable output string.
pub fn format_responses(responses: &[PartyResponse]) -> String {
    if responses.is_empty() {
        return "No personas responded.".to_string();
    }

    let mut output = String::new();
    for (i, resp) in responses.iter().enumerate() {
        if i > 0 {
            output.push('\n');
        }

        if let Some(err) = &resp.error {
            output.push_str(&format!(
                "{} **{}** [error]: {}\n",
                resp.persona.icon, resp.persona.name, err
            ));
        } else {
            output.push_str(&format!(
                "{} **{}** ({} round trips, {:.1}s):\n\n{}\n",
                resp.persona.icon,
                resp.persona.name,
                resp.round_trips,
                resp.elapsed.as_secs_f64(),
                resp.text
            ));
        }
    }

    output
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_persona(name: &str) -> Persona {
        Persona {
            skill_name: format!("{}-test", name.to_lowercase()),
            name: name.to_string(),
            title: format!("{} Title", name),
            icon: "🤖".to_string(),
            role: format!("{} role", name),
            identity: format!("{} identity", name),
            communication_style: "Professional".to_string(),
            principles: vec!["Be helpful".to_string()],
            body: format!("# {}\n\nBody.", name),
            when_to_use: String::new(),
            tools: None,
            skills: Vec::new(),
            ceiling: crate::risk::Capability::ReadOnly,
        }
    }

    // Equal budgets meant the outer timeout always won, so the inner one's round-attributed
    // message could never fire and the text produced before a stall was discarded with it.
    #[test]
    fn a_single_request_is_bounded_more_tightly_than_the_whole_persona() {
        let total = Duration::from_secs(120);
        let share = request_share(total);

        assert!(share < total, "{:?} must be under {:?}", share, total);
        assert_eq!(share, Duration::from_secs(96));
    }

    // A caller may pass a tiny timeout; the share must stay a duration a request can survive
    // rather than collapsing to zero and failing on the first poll.
    #[test]
    fn a_tiny_budget_still_leaves_a_usable_share() {
        assert_eq!(
            request_share(Duration::from_secs(1)),
            Duration::from_secs(1)
        );
        assert_eq!(request_share(Duration::ZERO), Duration::from_secs(1));
    }

    #[test]
    fn format_responses_with_no_responses() {
        let output = format_responses(&[]);
        assert_eq!(output, "No personas responded.");
    }

    #[test]
    fn format_responses_with_success() {
        let responses = vec![PartyResponse {
            persona: make_persona("Tyler"),
            text: "Design the contract.".to_string(),
            round_trips: 2,
            elapsed: Duration::from_secs_f64(1.5),
            error: None,
        }];

        let output = format_responses(&responses);
        assert!(output.contains("Tyler"));
        assert!(output.contains("Design the contract."));
        assert!(output.contains("2 round trips"));
    }

    #[test]
    fn format_responses_with_error() {
        let responses = vec![PartyResponse {
            persona: make_persona("Elliot"),
            text: String::new(),
            round_trips: 0,
            elapsed: Duration::ZERO,
            error: Some("Timed out".to_string()),
        }];

        let output = format_responses(&responses);
        assert!(output.contains("Elliot"));
        assert!(output.contains("error"));
        assert!(output.contains("Timed out"));
    }

    #[test]
    fn format_responses_multiple() {
        let responses = vec![
            PartyResponse {
                persona: make_persona("Tyler"),
                text: "I think we should use Soroban.".to_string(),
                round_trips: 1,
                elapsed: Duration::from_secs(1),
                error: None,
            },
            PartyResponse {
                persona: make_persona("Elliot"),
                text: "I agree, let me implement it.".to_string(),
                round_trips: 1,
                elapsed: Duration::from_secs(2),
                error: None,
            },
        ];

        let output = format_responses(&responses);
        assert!(output.contains("Tyler"));
        assert!(output.contains("Elliot"));
        assert!(output.contains("Soroban"));
        assert!(output.contains("implement"));
    }
}