use crate::{
    invocation::{SkillContext, SkillInput},
    outcome::SkillOutcome,
    skill::{Skill, SkillId},
};
use forge::budget::estimator::TokenEstimator;
use serde_json::json;
use weave::lenses::{
    clarity::ClarityLens, depth::DepthLens, focus::FocusLens, harmonic::HarmonicLens,
    narrow::NarrowLens, refract::RefractLens, wide::WideLens, Lens, LensContext,
};

/// Crucible: run compression benchmarks against sample inputs, return savings table.
pub struct CrucibleSkill;

impl Skill for CrucibleSkill {
    fn id(&self) -> SkillId {
        SkillId("crucible")
    }
    fn display_name(&self) -> &str {
        "Crucible"
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "samples": {
                    "type": "array",
                    "items": { "type": "string" },
                    "description": "Text samples to benchmark"
                },
                "lenses": {
                    "type": "array",
                    "items": { "type": "string" },
                    "description": "Lens names to benchmark (omit for all)"
                },
                "budget_tokens": { "type": "integer", "default": 2000 }
            },
            "required": ["samples"]
        })
    }

    fn output_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "results": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "lens": { "type": "string" },
                            "avg_savings_pct": { "type": "number" },
                            "avg_tokens_before": { "type": "number" },
                            "avg_tokens_after": { "type": "number" }
                        }
                    }
                }
            }
        })
    }

    fn execute(&self, input: SkillInput, _ctx: &SkillContext) -> anyhow::Result<SkillOutcome> {
        let samples = input.get_str_array("samples").unwrap_or_default();
        let requested = input.get_str_array("lenses").unwrap_or_default();
        let budget = input.get_u64("budget_tokens").unwrap_or(2000) as usize;

        type ApplyFn = Box<dyn Fn(&str, &LensContext) -> String>;
        let lenses_to_test: Vec<(&str, ApplyFn)> = vec![
            (
                "clarity",
                Box::new(|s: &str, c: &LensContext| ClarityLens.apply(s, c).content),
            ),
            (
                "narrow",
                Box::new(|s: &str, c: &LensContext| NarrowLens.apply(s, c).content),
            ),
            (
                "depth",
                Box::new(|s: &str, c: &LensContext| DepthLens.apply(s, c).content),
            ),
            (
                "focus",
                Box::new(|s: &str, c: &LensContext| FocusLens.apply(s, c).content),
            ),
            (
                "refract",
                Box::new(|s: &str, c: &LensContext| RefractLens.apply(s, c).content),
            ),
            (
                "wide",
                Box::new(|s: &str, c: &LensContext| WideLens.apply(s, c).content),
            ),
            (
                "harmonic",
                Box::new(|s: &str, c: &LensContext| HarmonicLens.apply(s, c).content),
            ),
        ];

        let mut results = Vec::new();

        for (name, apply_fn) in lenses_to_test {
            if !requested.is_empty() && !requested.iter().any(|r| r == name) {
                continue;
            }

            let ctx = LensContext::new(budget);
            let mut total_before = 0usize;
            let mut total_after = 0usize;

            for sample in &samples {
                let before = TokenEstimator::count_nonblocking(sample);
                let after_str = apply_fn(sample, &ctx);
                let after = TokenEstimator::count_nonblocking(&after_str);
                total_before += before;
                total_after += after;
            }

            let n = samples.len().max(1);
            let avg_before = total_before as f64 / n as f64;
            let avg_after = total_after as f64 / n as f64;
            let avg_savings = TokenEstimator::savings_pct(total_before, total_after);

            results.push(json!({
                "lens": name,
                "avg_savings_pct": (avg_savings * 10.0).round() / 10.0,
                "avg_tokens_before": (avg_before * 10.0).round() / 10.0,
                "avg_tokens_after": (avg_after * 10.0).round() / 10.0
            }));
        }

        // Sort by savings descending
        results.sort_by(|a, b| {
            b["avg_savings_pct"]
                .as_f64()
                .unwrap_or(0.0)
                .partial_cmp(&a["avg_savings_pct"].as_f64().unwrap_or(0.0))
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        Ok(SkillOutcome::success(
            "crucible",
            json!({ "results": results }),
        ))
    }
}
