// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::time::Duration;

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use psmatcher::{
    event_types::{AttributeValue, DefaultEvent},
    ElementaryPredicate, ElementaryTest, Matcher, Subscription,
};
use rand::{rngs::StdRng, Rng, SeedableRng};

// Generate a random event with given attributes
fn generate_random_event(rng: &mut StdRng, id: usize) -> DefaultEvent {
    let topic = match rng.gen_range(0..5) {
        0 => "news",
        1 => "weather",
        2 => "sports",
        3 => "finance",
        _ => "tech",
    };

    let category = match rng.gen_range(0..10) {
        0 => "breaking",
        1 => "local",
        2 => "global",
        3 => "technology",
        4 => "science",
        5 => "health",
        6 => "politics",
        7 => "business",
        8 => "entertainment",
        _ => "other",
    };

    let region = match rng.gen_range(0..5) {
        0 => "north",
        1 => "south",
        2 => "east",
        3 => "west",
        _ => "central",
    };

    let priority = rng.gen_range(1..10);
    let value = rng.gen_range(0.0..100.0);

    DefaultEvent::new(topic)
        .with_attribute("id", AttributeValue::Integer(id as i64))
        .with_attribute("category", AttributeValue::String(category.to_string()))
        .with_attribute("region", AttributeValue::String(region.to_string()))
        .with_attribute("priority", AttributeValue::Integer(priority))
        .with_attribute("value", AttributeValue::Float(value))
}

// Generate a set of random events
fn generate_events(count: usize) -> Vec<DefaultEvent> {
    let mut rng = StdRng::seed_from_u64(42); // Fixed seed for reproducibility
    (0..count)
        .map(|i| generate_random_event(&mut rng, i))
        .collect()
}

// Generate a random subscription with 2 predicates
fn generate_random_subscription(rng: &mut StdRng, id: usize) -> Subscription<String> {
    let sub_id = format!("sub-{id}");

    // First predicate is always on topic
    let topic_predicate = match rng.gen_range(0..5) {
        0 => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "topic".to_string(),
                value: AttributeValue::String("news".to_string()),
            },
            false,
        ),
        1 => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "topic".to_string(),
                value: AttributeValue::String("weather".to_string()),
            },
            false,
        ),
        2 => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "topic".to_string(),
                value: AttributeValue::String("sports".to_string()),
            },
            false,
        ),
        3 => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "topic".to_string(),
                value: AttributeValue::String("finance".to_string()),
            },
            false,
        ),
        _ => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "topic".to_string(),
                value: AttributeValue::String("tech".to_string()),
            },
            false,
        ),
    };

    // Second predicate varies
    let second_predicate = match rng.gen_range(0..5) {
        0 => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "category".to_string(),
                value: AttributeValue::String("breaking".to_string()),
            },
            false,
        ),
        1 => ElementaryPredicate::new(
            ElementaryTest::GreaterThan {
                attribute: "priority".to_string(),
                value: 5.0,
            },
            false,
        ),
        2 => ElementaryPredicate::new(
            ElementaryTest::LessThan {
                attribute: "value".to_string(),
                value: 50.0,
            },
            false,
        ),
        3 => ElementaryPredicate::new(
            ElementaryTest::Equals {
                attribute: "region".to_string(),
                value: AttributeValue::String("north".to_string()),
            },
            false,
        ),
        _ => ElementaryPredicate::new(
            ElementaryTest::regex("category".to_string(), "tech.*|science").unwrap(),
            false,
        ),
    };

    Subscription::new(sub_id)
        .with_predicate(topic_predicate)
        .with_predicate(second_predicate)
}

// Generate a set of random subscriptions
fn generate_subscriptions(count: usize) -> Vec<Subscription<String>> {
    let mut rng = StdRng::seed_from_u64(123); // Fixed seed for reproducibility
    (0..count)
        .map(|i| generate_random_subscription(&mut rng, i))
        .collect()
}

// Setup the matcher with subscriptions
fn setup_matcher(subscription_count: usize) -> Matcher<String> {
    let subscriptions = generate_subscriptions(subscription_count);
    let mut matcher = Matcher::new();

    for subscription in subscriptions {
        matcher.add_subscription(subscription);
    }

    matcher
}

// Benchmark function
fn bench_matcher(c: &mut Criterion) {
    let subscription_count = 100;
    let matcher = setup_matcher(subscription_count);

    let mut group = c.benchmark_group("event_matching");
    group.measurement_time(Duration::from_secs(10));
    group.sample_size(50);

    for event_count in [10, 100, 1000].iter() {
        let events = generate_events(*event_count);

        group.bench_with_input(
            BenchmarkId::from_parameter(event_count),
            &events,
            |b, events| {
                b.iter(|| {
                    for event in events {
                        let _ = matcher.match_event(event).unwrap();
                    }
                })
            },
        );
    }

    group.finish();
}

// Benchmark function for individual event matching
fn bench_individual_event_matching(c: &mut Criterion) {
    let subscription_count = 100;
    let matcher = setup_matcher(subscription_count);

    let mut group = c.benchmark_group("individual_event_matching");
    group.measurement_time(Duration::from_secs(10));
    group.sample_size(100);

    // Generate a few representative events
    let mut rng = StdRng::seed_from_u64(42);
    let event1 = generate_random_event(&mut rng, 1); // A typical event
    let event2 = DefaultEvent::new("news")
        .with_attribute("category", AttributeValue::String("breaking".to_string()))
        .with_attribute("priority".to_string(), AttributeValue::Integer(9)); // High match event
    let event3 = DefaultEvent::new("custom")
        .with_attribute("category", AttributeValue::String("rare".to_string())); // Low match event

    group.bench_function("typical_event", |b| {
        b.iter(|| matcher.match_event(&event1).unwrap())
    });

    group.bench_function("high_match_event", |b| {
        b.iter(|| matcher.match_event(&event2).unwrap())
    });

    group.bench_function("low_match_event", |b| {
        b.iter(|| matcher.match_event(&event3).unwrap())
    });

    group.finish();
}

// Benchmark function for varying subscription counts
fn bench_subscription_scaling(c: &mut Criterion) {
    let mut group = c.benchmark_group("subscription_scaling");
    group.measurement_time(Duration::from_secs(10));
    group.sample_size(50);

    let events = generate_events(10); // Fixed number of events

    for sub_count in [10, 100, 1000].iter() {
        let matcher = setup_matcher(*sub_count);

        group.bench_with_input(
            BenchmarkId::from_parameter(sub_count),
            &events,
            |b, events| {
                b.iter(|| {
                    for event in events {
                        let _ = matcher.match_event(event).unwrap();
                    }
                })
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_matcher,
    bench_individual_event_matching,
    bench_subscription_scaling
);
criterion_main!(benches);
