use super::{group_by_module, Event, EventListener};
use aiken_lang::{
    ast::OnTestFailure,
    expr::UntypedExpr,
    format::Formatter,
    test_framework::{AssertionStyleOptions, PropertyTestResult, TestResult, UnitTestResult},
};
use serde_json::json;

#[derive(Debug, Default, Clone, Copy)]
pub struct Json;

impl EventListener for Json {
    fn handle_event(&self, event: Event) {
        match event {
            Event::FinishedTests { seed, tests, .. } => {
                let total = tests.len();
                let passed = tests.iter().filter(|t| t.is_success()).count();
                let failed = total - passed;

                let json_output = serde_json::json!({
                    "seed": seed,
                    "summary": json!({
                        "total": total,
                        "passed": passed,
                        "failed": failed,
                        "kind": json!({
                            "unit": count_unit_tests(tests.iter()),
                            "property": count_property_tests(tests.iter()),
                        })
                    }),
                    "modules": group_by_module(&tests).iter().map(|(module, results)| {
                        serde_json::json!({
                            "name": module,
                            "summary": fmt_test_summary_json(results),
                            "tests": results.iter().map(|r| fmt_test_json(r)).collect::<Vec<_>>(),
                        })
                    }).collect::<Vec<_>>(),
                });
                println!("{}", serde_json::to_string_pretty(&json_output).unwrap());
            }
            _ => super::Terminal.handle_event(event),
        }
    }
}

fn fmt_test_json(result: &TestResult<UntypedExpr, UntypedExpr>) -> serde_json::Value {
    let mut test = json!({
        "title": result.title(),
        "status": if result.is_success() { "pass" } else { "fail" },
    });

    match result {
        TestResult::UnitTestResult(UnitTestResult {
            spent_budget,
            assertion,
            test: unit_test,
            ..
        }) => {
            test["execution_units"] = json!({
                "mem": spent_budget.mem,
                "cpu": spent_budget.cpu,
            });
            if !result.is_success() {
                if let Some(assertion) = assertion {
                    test["assertion"] = json!({
                        "message": assertion.to_string(false, &AssertionStyleOptions::new(None)),
                        "on_failure": match unit_test.on_test_failure {
                            OnTestFailure::FailImmediately => "fail_immediately" ,
                            OnTestFailure::SucceedEventually => "succeed_eventually" ,
                            OnTestFailure::SucceedImmediately => "succeed_immediately",
                        }
                    });
                }
            }
        }
        TestResult::PropertyTestResult(PropertyTestResult {
            iterations,
            labels,
            counterexample,
            ..
        }) => {
            test["iterations"] = json!(iterations);
            test["labels"] = json!(labels);
            test["counterexample"] = match counterexample {
                Ok(Some(expr)) => json!(Formatter::new().expr(expr, false).to_pretty_string(60)),
                Ok(None) => json!(null),
                Err(err) => json!({"error": err.to_string()}),
            };
        }
    }

    if !result.traces().is_empty() {
        test["traces"] = json!(result.traces());
    }

    test
}

fn fmt_test_summary_json(tests: &[&TestResult<UntypedExpr, UntypedExpr>]) -> serde_json::Value {
    let total = tests.len();
    let passed = tests.iter().filter(|t| t.is_success()).count();
    let failed = total - passed;

    json!({
        "total": total,
        "passed": passed,
        "failed": failed,
        "kind": json!({
            "unit": count_unit_tests(tests.iter().copied()),
            "property": count_property_tests(tests.iter().copied())
        })
    })
}

fn count_unit_tests<'a, I>(tests: I) -> usize
where
    I: Iterator<Item = &'a TestResult<UntypedExpr, UntypedExpr>>,
{
    tests
        .filter(|t| matches!(t, TestResult::UnitTestResult { .. }))
        .count()
}

fn count_property_tests<'a, I>(tests: I) -> usize
where
    I: Iterator<Item = &'a TestResult<UntypedExpr, UntypedExpr>>,
{
    tests
        .filter(|t| matches!(t, TestResult::PropertyTestResult { .. }))
        .count()
}