swarm-engine-llm 0.1.6

//! Strategy Advisor - LLM による探索戦略アドバイス
//!
//! 探索の状態に基づいて最適な Selection 戦略を推奨する。
//!
//! # 設計
//!
//! ```text
//! StrategyContext (探索状態)
//!         │
//!         ▼
//! StrategyPromptBuilder.build()
//!         │
//!         ▼
//! LlmDecider.call_raw() ──────── 同期ブロッキング (~100ms)
//!         │
//!         ▼
//! StrategyResponseParser.parse()
//!         │
//!         ▼
//! StrategyAdvice (推奨戦略)
//! ```
//!
//! # 使用例
//!
//! ```ignore
//! use swarm_engine_llm::strategy_advisor::{LlmStrategyAdvisor, StrategyContext};
//! use swarm_engine_core::exploration::SelectionKind;
//!
//! let advisor = LlmStrategyAdvisor::new(decider, runtime);
//! let context = StrategyContext::new(15, 47, 0.23, SelectionKind::Ucb1);
//! let advice = advisor.advise(&context)?;
//! ```

use std::sync::Arc;

use fuzzy_parser::distance::{find_closest, Algorithm};
use fuzzy_parser::{repair_object_fields, sanitize_json, ObjectSchema};

// Core 層から型をインポート
pub use swarm_engine_core::exploration::{
    SelectionKind, StrategyAdvice, StrategyAdviceError, StrategyAdvisor, StrategyContext,
};

use crate::decider::{LlmDecider, LlmError, LoraConfig};
use crate::json_prompt::strategy_selection_template;

// ============================================================================
// SelectionKind 拡張 - 文字列からの fuzzy パース
// ============================================================================

/// SelectionKind の文字列からの fuzzy パース
pub fn parse_selection_kind_fuzzy(s: &str) -> Option<SelectionKind> {
    // 完全一致（大文字小文字無視）
    let upper = s.to_uppercase();
    match upper.as_str() {
        "FIFO" => return Some(SelectionKind::Fifo),
        "UCB1" => return Some(SelectionKind::Ucb1),
        "GREEDY" => return Some(SelectionKind::Greedy),
        "THOMPSON" => return Some(SelectionKind::Thompson),
        _ => {}
    }

    // Fuzzy match
    let candidates = ["FIFO", "UCB1", "Greedy", "Thompson"];
    if let Some(m) = find_closest(s, candidates, 0.6, Algorithm::JaroWinkler) {
        match m.candidate.as_str() {
            "FIFO" => Some(SelectionKind::Fifo),
            "UCB1" => Some(SelectionKind::Ucb1),
            "Greedy" => Some(SelectionKind::Greedy),
            "Thompson" => Some(SelectionKind::Thompson),
            _ => None,
        }
    } else {
        None
    }
}

// ============================================================================
// StrategyAdviceError 拡張
// ============================================================================

impl From<LlmError> for StrategyAdviceError {
    fn from(e: LlmError) -> Self {
        Self::LlmError(e.message().to_string())
    }
}

// ============================================================================
// StrategyPromptBuilder - プロンプト生成
// ============================================================================

/// 戦略アドバイス用プロンプトビルダー
#[derive(Debug, Clone, Default)]
pub struct StrategyPromptBuilder;

impl StrategyPromptBuilder {
    /// 新規作成
    pub fn new() -> Self {
        Self
    }

    /// プロンプトを生成
    pub fn build(&self, ctx: &StrategyContext) -> String {
        let depth_info = ctx
            .avg_depth
            .map(|d| format!(", depth={:.1}", d))
            .unwrap_or_default();

        // コンテキスト情報を構築
        let content = format!(
            "Strategies: FIFO, UCB1, Greedy, Thompson\n\
             Guidelines: visits<20→UCB1, failure>30%→Thompson, established+low failure→Greedy\n\
             User: frontier={}, visits={}, failure={:.0}%{}, current={}",
            ctx.frontier_count,
            ctx.total_visits,
            ctx.failure_rate * 100.0,
            depth_info,
            ctx.current_strategy,
        );

        // 共通テンプレートを使用
        strategy_selection_template().build(&content)
    }
}

// ============================================================================
// StrategyResponseParser - レスポンスパース
// ============================================================================

/// 戦略レスポンス用 ObjectSchema
const STRATEGY_FIELDS: ObjectSchema =
    ObjectSchema::new(&["strategy", "change", "confidence", "reason"]);

/// 戦略アドバイス用レスポンスパーサー
#[derive(Debug, Clone, Default)]
pub struct StrategyResponseParser;

impl StrategyResponseParser {
    /// 新規作成
    pub fn new() -> Self {
        Self
    }

    /// レスポンスをパース
    pub fn parse(&self, response: &str) -> Result<StrategyAdvice, StrategyAdviceError> {
        // JSON を抽出
        let json_str = self.extract_json(response)?;

        // 構文修復
        let sanitized = sanitize_json(&json_str);
        tracing::debug!(sanitized = %sanitized, "Sanitized strategy JSON");

        // パース
        self.parse_json(&sanitized)
    }

    /// JSON を抽出（```json ブロック対応、自然言語フォールバック）
    fn extract_json(&self, text: &str) -> Result<String, StrategyAdviceError> {
        // ```json ... ``` ブロックを探す
        if let Some(start) = text.find("```json") {
            let content_start = start + 7;
            let remaining = &text[content_start..];
            if let Some(end) = remaining.find("```") {
                let json = remaining[..end].trim();
                if !json.is_empty() {
                    return Ok(json.to_string());
                }
            }
        }

        // { ... } を探す（バランスを取って）
        if let Some(json) = self.extract_balanced_json(text) {
            return Ok(json);
        }

        // フォールバック: 自然言語から戦略キーワードを抽出して JSON 生成
        if let Some(json) = self.extract_from_natural_language(text) {
            tracing::debug!(fallback_json = %json, "Extracted strategy from natural language");
            return Ok(json);
        }

        Err(StrategyAdviceError::ParseError(format!(
            "No JSON found in response: {}",
            text
        )))
    }

    /// 自然言語から戦略キーワードを抽出して JSON 生成
    fn extract_from_natural_language(&self, text: &str) -> Option<String> {
        let text_upper = text.to_uppercase();

        // 優先度順に戦略を検索（推奨を示す文脈を考慮）
        let recommend_patterns = ["RECOMMEND", "SUGGEST", "USE ", "PREFER", "OPTIMAL", "BEST"];
        let strategies = [
            ("THOMPSON", "Thompson"),
            ("UCB1", "UCB1"),
            ("UCB", "UCB1"),
            ("GREEDY", "Greedy"),
            ("FIFO", "FIFO"),
        ];

        // まず推奨文脈付きの戦略を探す
        for pattern in &recommend_patterns {
            if let Some(pos) = text_upper.find(pattern) {
                // パターン後の50文字以内で戦略を探す
                let search_range = &text_upper[pos..std::cmp::min(pos + 50, text_upper.len())];
                for (keyword, strategy) in &strategies {
                    if search_range.contains(keyword) {
                        return Some(format!(
                            r#"{{"strategy":"{}","change":true,"confidence":0.6,"reason":"Extracted from natural language response"}}"#,
                            strategy
                        ));
                    }
                }
            }
        }

        // 推奨文脈がなければ、最初に出現した戦略を使用
        let mut first_match: Option<(usize, &str)> = None;
        for (keyword, strategy) in &strategies {
            if let Some(pos) = text_upper.find(keyword) {
                if first_match.is_none() || pos < first_match.unwrap().0 {
                    first_match = Some((pos, strategy));
                }
            }
        }

        first_match.map(|(_, strategy)| {
            format!(
                r#"{{"strategy":"{}","change":false,"confidence":0.5,"reason":"Extracted from natural language response"}}"#,
                strategy
            )
        })
    }

    /// バランスの取れた JSON を抽出
    fn extract_balanced_json(&self, text: &str) -> Option<String> {
        let start = text.find('{')?;
        let chars: Vec<char> = text[start..].chars().collect();
        let mut depth = 0;
        let mut in_string = false;
        let mut escape_next = false;

        for (i, &ch) in chars.iter().enumerate() {
            if escape_next {
                escape_next = false;
                continue;
            }

            match ch {
                '\\' if in_string => escape_next = true,
                '"' => in_string = !in_string,
                '{' if !in_string => depth += 1,
                '}' if !in_string => {
                    depth -= 1;
                    if depth == 0 {
                        return Some(chars[..=i].iter().collect());
                    }
                }
                _ => {}
            }
        }

        None
    }

    /// JSON をパース（fuzzy repair 対応）
    fn parse_json(&self, json: &str) -> Result<StrategyAdvice, StrategyAdviceError> {
        let mut parsed: serde_json::Value = serde_json::from_str(json)
            .map_err(|e| StrategyAdviceError::ParseError(format!("JSON parse error: {}", e)))?;

        // フィールド名の typo 修復
        if let Some(obj) = parsed.as_object_mut() {
            let corrections = repair_object_fields(obj, &STRATEGY_FIELDS, "$", &Default::default());
            if !corrections.is_empty() {
                tracing::debug!(
                    corrections = ?corrections.iter().map(|c| format!("{} → {}", c.original, c.corrected)).collect::<Vec<_>>(),
                    "Fuzzy repaired strategy field names"
                );
            }
        }

        // strategy フィールドをパース（fuzzy repair 対応）
        let strategy_str = parsed["strategy"]
            .as_str()
            .ok_or_else(|| StrategyAdviceError::ParseError("Missing 'strategy' field".into()))?;

        let recommended = parse_selection_kind_fuzzy(strategy_str).ok_or_else(|| {
            StrategyAdviceError::ParseError(format!("Unknown strategy: {}", strategy_str))
        })?;

        let should_change = parsed["change"].as_bool().unwrap_or(false);
        let confidence = parsed["confidence"].as_f64().unwrap_or(0.5).clamp(0.0, 1.0);
        let reason = parsed["reason"]
            .as_str()
            .unwrap_or("No reason provided")
            .to_string();

        Ok(StrategyAdvice {
            recommended,
            should_change,
            reason,
            confidence,
        })
    }
}

// ============================================================================
// LlmStrategyAdvisor - LLM ベースの実装
// ============================================================================

/// LLM ベースの戦略アドバイザー
pub struct LlmStrategyAdvisor {
    decider: Arc<dyn LlmDecider>,
    runtime: tokio::runtime::Handle,
    prompt_builder: StrategyPromptBuilder,
    response_parser: StrategyResponseParser,
    /// 信頼度閾値（これ以下のアドバイスは変更しない）
    confidence_threshold: f64,
    /// LoRA 設定（None の場合はベースモデルのみ）
    lora: Option<LoraConfig>,
}

impl LlmStrategyAdvisor {
    /// 新しい LlmStrategyAdvisor を作成
    pub fn new(decider: Arc<dyn LlmDecider>, runtime: tokio::runtime::Handle) -> Self {
        Self {
            decider,
            runtime,
            prompt_builder: StrategyPromptBuilder::new(),
            response_parser: StrategyResponseParser::new(),
            confidence_threshold: 0.6,
            lora: None,
        }
    }

    /// 信頼度閾値を設定
    pub fn with_confidence_threshold(mut self, threshold: f64) -> Self {
        self.confidence_threshold = threshold.clamp(0.0, 1.0);
        self
    }

    /// 信頼度閾値を取得
    pub fn confidence_threshold(&self) -> f64 {
        self.confidence_threshold
    }

    /// LoRA 設定を設定
    pub fn with_lora(mut self, lora: LoraConfig) -> Self {
        self.lora = Some(lora);
        self
    }

    /// LoRA 設定を取得
    pub fn lora(&self) -> Option<&LoraConfig> {
        self.lora.as_ref()
    }
}

impl StrategyAdvisor for LlmStrategyAdvisor {
    fn advise(&self, context: &StrategyContext) -> Result<StrategyAdvice, StrategyAdviceError> {
        // プロンプト生成
        let prompt = self.prompt_builder.build(context);
        tracing::debug!(prompt = %prompt, "Strategy advisor prompt");

        // 同期ブロッキング呼び出し（~100ms 想定）
        let lora = self.lora.as_ref();
        let response = self
            .runtime
            .block_on(async { self.decider.call_raw(&prompt, lora).await })?;

        tracing::debug!(response = %response, "Strategy advisor raw response");

        // レスポンスパース
        let mut advice = self.response_parser.parse(&response)?;

        // 信頼度が閾値以下なら変更しない
        if advice.confidence < self.confidence_threshold {
            tracing::debug!(
                confidence = advice.confidence,
                threshold = self.confidence_threshold,
                "Low confidence, not changing strategy"
            );
            advice.should_change = false;
            advice.reason = format!(
                "Low confidence ({:.2} < {:.2}): {}",
                advice.confidence, self.confidence_threshold, advice.reason
            );
        }

        tracing::info!(
            recommended = %advice.recommended,
            should_change = advice.should_change,
            confidence = advice.confidence,
            reason = %advice.reason,
            "Strategy advice"
        );

        Ok(advice)
    }

    fn name(&self) -> &str {
        "LlmStrategyAdvisor"
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    // ========================================================================
    // SelectionKind Tests
    // ========================================================================

    #[test]
    fn test_selection_kind_display() {
        assert_eq!(SelectionKind::Fifo.to_string(), "FIFO");
        assert_eq!(SelectionKind::Ucb1.to_string(), "UCB1");
        assert_eq!(SelectionKind::Greedy.to_string(), "Greedy");
        assert_eq!(SelectionKind::Thompson.to_string(), "Thompson");
    }

    #[test]
    fn test_selection_kind_from_str_exact() {
        assert_eq!(
            parse_selection_kind_fuzzy("FIFO"),
            Some(SelectionKind::Fifo)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("UCB1"),
            Some(SelectionKind::Ucb1)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("Greedy"),
            Some(SelectionKind::Greedy)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("Thompson"),
            Some(SelectionKind::Thompson)
        );
    }

    #[test]
    fn test_selection_kind_from_str_case_insensitive() {
        assert_eq!(
            parse_selection_kind_fuzzy("fifo"),
            Some(SelectionKind::Fifo)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("ucb1"),
            Some(SelectionKind::Ucb1)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("GREEDY"),
            Some(SelectionKind::Greedy)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("THOMPSON"),
            Some(SelectionKind::Thompson)
        );
    }

    #[test]
    fn test_selection_kind_from_str_fuzzy() {
        // Typos should be repaired
        assert_eq!(
            parse_selection_kind_fuzzy("Thomspon"),
            Some(SelectionKind::Thompson)
        );
        assert_eq!(
            parse_selection_kind_fuzzy("Gredy"),
            Some(SelectionKind::Greedy)
        );
    }

    #[test]
    fn test_selection_kind_from_str_invalid() {
        assert_eq!(parse_selection_kind_fuzzy("Unknown"), None);
        assert_eq!(parse_selection_kind_fuzzy("Random"), None);
    }

    // ========================================================================
    // StrategyContext Tests
    // ========================================================================

    #[test]
    fn test_strategy_context_new() {
        let ctx = StrategyContext::new(15, 47, 0.23, SelectionKind::Ucb1);
        assert_eq!(ctx.frontier_count, 15);
        assert_eq!(ctx.total_visits, 47);
        assert!((ctx.failure_rate - 0.23).abs() < 0.001);
        assert!((ctx.success_rate - 0.77).abs() < 0.001);
        assert_eq!(ctx.current_strategy, SelectionKind::Ucb1);
        assert!(ctx.avg_depth.is_none());
    }

    #[test]
    fn test_strategy_context_with_depth() {
        let ctx = StrategyContext::new(10, 100, 0.1, SelectionKind::Greedy).with_avg_depth(3.5);
        assert_eq!(ctx.avg_depth, Some(3.5));
    }

    // ========================================================================
    // StrategyAdvice Tests
    // ========================================================================

    #[test]
    fn test_strategy_advice_no_change() {
        let advice = StrategyAdvice::no_change(SelectionKind::Ucb1, "Exploration phase");
        assert_eq!(advice.recommended, SelectionKind::Ucb1);
        assert!(!advice.should_change);
        assert_eq!(advice.reason, "Exploration phase");
        assert!((advice.confidence - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_strategy_advice_change_to() {
        let advice = StrategyAdvice::change_to(SelectionKind::Greedy, "Patterns established", 0.85);
        assert_eq!(advice.recommended, SelectionKind::Greedy);
        assert!(advice.should_change);
        assert_eq!(advice.reason, "Patterns established");
        assert!((advice.confidence - 0.85).abs() < 0.001);
    }

    // ========================================================================
    // StrategyPromptBuilder Tests
    // ========================================================================

    #[test]
    fn test_prompt_builder_basic() {
        let builder = StrategyPromptBuilder::new();
        let ctx = StrategyContext::new(15, 47, 0.23, SelectionKind::Ucb1);
        let prompt = builder.build(&ctx);

        // Few-shot example format
        assert!(prompt.contains("Example interaction:"));
        assert!(prompt.contains("Your JSON:"));

        // Context values
        assert!(prompt.contains("frontier=15"));
        assert!(prompt.contains("visits=47"));
        assert!(prompt.contains("failure=23%"));
        assert!(prompt.contains("current=UCB1"));

        // Strategy names
        assert!(prompt.contains("FIFO"));
        assert!(prompt.contains("Greedy"));
        assert!(prompt.contains("Thompson"));
    }

    #[test]
    fn test_prompt_builder_with_depth() {
        let builder = StrategyPromptBuilder::new();
        let ctx = StrategyContext::new(10, 100, 0.1, SelectionKind::Greedy).with_avg_depth(3.5);
        let prompt = builder.build(&ctx);

        assert!(prompt.contains("depth=3.5"));
    }

    // ========================================================================
    // StrategyResponseParser Tests
    // ========================================================================

    #[test]
    fn test_parse_valid_json() {
        let parser = StrategyResponseParser::new();
        let response = r#"{"strategy": "Greedy", "change": true, "confidence": 0.85, "reason": "Low failure rate"}"#;
        let advice = parser.parse(response).unwrap();

        assert_eq!(advice.recommended, SelectionKind::Greedy);
        assert!(advice.should_change);
        assert!((advice.confidence - 0.85).abs() < 0.001);
        assert_eq!(advice.reason, "Low failure rate");
    }

    #[test]
    fn test_parse_json_with_prefix() {
        let parser = StrategyResponseParser::new();
        let response = r#"Based on the analysis: {"strategy": "Thompson", "change": true, "confidence": 0.7, "reason": "High variance"}"#;
        let advice = parser.parse(response).unwrap();

        assert_eq!(advice.recommended, SelectionKind::Thompson);
    }

    #[test]
    fn test_parse_json_markdown_block() {
        let parser = StrategyResponseParser::new();
        let response = r#"```json
{"strategy": "UCB1", "change": false, "confidence": 0.9, "reason": "Still exploring"}
```"#;
        let advice = parser.parse(response).unwrap();

        assert_eq!(advice.recommended, SelectionKind::Ucb1);
        assert!(!advice.should_change);
    }

    #[test]
    fn test_parse_json_typo_repair() {
        let parser = StrategyResponseParser::new();
        // "straegy" typo should be repaired to "strategy"
        let response =
            r#"{"straegy": "Greedy", "change": true, "confidnce": 0.8, "reason": "test"}"#;
        let advice = parser.parse(response).unwrap();

        assert_eq!(advice.recommended, SelectionKind::Greedy);
    }

    #[test]
    fn test_parse_json_strategy_typo() {
        let parser = StrategyResponseParser::new();
        // "Thomspon" typo should be repaired to "Thompson"
        let response =
            r#"{"strategy": "Thomspon", "change": true, "confidence": 0.75, "reason": "variance"}"#;
        let advice = parser.parse(response).unwrap();

        assert_eq!(advice.recommended, SelectionKind::Thompson);
    }

    #[test]
    fn test_parse_json_defaults() {
        let parser = StrategyResponseParser::new();
        // Missing change and confidence should use defaults
        let response = r#"{"strategy": "FIFO", "reason": "simple"}"#;
        let advice = parser.parse(response).unwrap();

        assert_eq!(advice.recommended, SelectionKind::Fifo);
        assert!(!advice.should_change); // default false
        assert!((advice.confidence - 0.5).abs() < 0.001); // default 0.5
    }

    #[test]
    fn test_parse_json_missing_strategy() {
        let parser = StrategyResponseParser::new();
        let response = r#"{"change": true, "confidence": 0.8}"#;
        let result = parser.parse(response);

        assert!(result.is_err());
        assert!(matches!(result, Err(StrategyAdviceError::ParseError(_))));
    }

    #[test]
    fn test_parse_no_json() {
        let parser = StrategyResponseParser::new();
        let response = "This is just plain text without any JSON.";
        let result = parser.parse(response);

        assert!(result.is_err());
        assert!(matches!(result, Err(StrategyAdviceError::ParseError(_))));
    }

    #[test]
    fn test_parse_confidence_clamping() {
        let parser = StrategyResponseParser::new();
        // Confidence > 1.0 should be clamped
        let response =
            r#"{"strategy": "Greedy", "change": true, "confidence": 1.5, "reason": "test"}"#;
        let advice = parser.parse(response).unwrap();
        assert!((advice.confidence - 1.0).abs() < 0.001);

        // Confidence < 0.0 should be clamped
        let response =
            r#"{"strategy": "Greedy", "change": true, "confidence": -0.5, "reason": "test"}"#;
        let advice = parser.parse(response).unwrap();
        assert!((advice.confidence - 0.0).abs() < 0.001);
    }
}