use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3::Client;
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use lazy_static::lazy_static;
use log::info;
use rayon::prelude::*;
use redis::{RedisError, RedisResult};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use std::env;
use std::io::Read;
use std::io::Write;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use tokio::io::AsyncReadExt;
use warp::Filter;

struct RedisClient {
    client: redis::Client,
}

enum AggregationType {
    Total,
    ByProduct,
    ByCountry,
    ByFeed,
}

enum Resolution {
    Low,
    Full,
}

// TODO: use async and connection pool
impl RedisClient {
    fn new() -> Self {
        let client = redis::Client::open(
            env::var("REDIS_URL")
                .unwrap_or_else(|_| "redis://localhost:6379".to_string())
                .as_str(),
        )
        .expect("Failed to create Redis client");
        RedisClient { client }
    }

    fn get_connection(&self) -> redis::Connection {
        self.client
            .get_connection()
            .expect("Failed to get Redis connection")
    }
}

lazy_static! {
    static ref REDIS_CLIENT: RedisClient = RedisClient::new();
}

// having an isrc (e.g. USRC15701155), and start_date (e.g. 2014-01-01), and end_date (e.g. 2023-09-30),
// generate an array of keys in format of "USRC15701155_2014_1"
async fn generate_keys_from_isrc_and_start_date_and_end_date(
    isrc: String,
    start_date: String,
    end_date: String,
) -> Vec<String> {
    let mut keys: Vec<String> = Vec::with_capacity(150);
    let start_date_vec: Vec<&str> = start_date.split('-').collect();
    let end_date_vec: Vec<&str> = end_date.split('-').collect();
    let start_year: i32 = start_date_vec[0].parse().unwrap_or(0);
    let end_year: i32 = end_date_vec[0].parse().unwrap_or(0);
    let orig_start_month: i32 = start_date_vec
        .get(1)
        .and_then(|s| s.parse().ok())
        .unwrap_or(1);
    let orig_end_month: i32 = end_date_vec
        .get(1)
        .and_then(|s| s.parse().ok())
        .unwrap_or(12);
    for year in start_year..=end_year {
        let start_month = if year == start_year {
            orig_start_month
        } else {
            1
        };
        let end_month = if year == end_year { orig_end_month } else { 12 };
        for month in start_month..=end_month {
            let key = format!("{}_{}_{}", isrc, year, month);
            keys.push(key);
        }
    }
    keys
}

#[derive(Debug, Deserialize)]
struct Item {
    product_id: i64,
    distributor: String,
    feed_id: i64,
    country_code: String,
    download_activity_date: String,
    streams: i64,
    streams_with_skips: i64,
    skips: i64,
    streams_passive: i64,
    streams_active: i64,
    streams_collection: i64,
    streams_unknown: i64,
    saves: i64,
    sub_type_subscription: i64,
    sub_type_adsupported: i64,
    sub_type_midtier: i64,
}

#[derive(Debug, Serialize)]
struct ItemTotal {
    download_activity_date: String,
    streams: i64,
    streams_with_skips: i64,
    skips: i64,
    streams_passive: i64,
    streams_active: i64,
    streams_collection: i64,
    streams_unknown: i64,
    saves: i64,
    sub_type_subscription: i64,
    sub_type_adsupported: i64,
    sub_type_midtier: i64,
}

async fn decode_and_flatten(items_from_redis: RedisResult<Vec<Vec<u8>>>) -> Vec<Item> {
    let items_from_redis = items_from_redis.unwrap();
    let final_results: Arc<Mutex<Vec<Item>>> =
        Arc::new(Mutex::new(Vec::with_capacity(items_from_redis.len() * 31)));

    items_from_redis.par_iter().for_each(|gzipped_array| {
        let mut decoder = GzDecoder::new(&gzipped_array[..]);
        let mut buffer_string = String::with_capacity(1024 * 1024 * 5);
        decoder.read_to_string(&mut buffer_string).unwrap();
        let items: Vec<serde_json::Value> = serde_json::from_str(&buffer_string).unwrap();
        let mut final_results = final_results.lock().unwrap();
        for values in items {
            let item = Item {
                product_id: values[0].as_i64().unwrap(),
                distributor: values[1].as_str().unwrap().to_string(),
                feed_id: values[2].as_i64().unwrap(),
                country_code: values[3].as_str().unwrap().to_string(),
                download_activity_date: values[4].as_str().unwrap().to_string(),
                streams: values[5].as_i64().unwrap(),
                streams_with_skips: values[6].as_i64().unwrap(),
                skips: values[7].as_i64().unwrap(),
                streams_passive: values[8].as_i64().unwrap(),
                streams_active: values[9].as_i64().unwrap(),
                streams_collection: values[10].as_i64().unwrap(),
                streams_unknown: values[11].as_i64().unwrap(),
                saves: values[12].as_i64().unwrap(),
                sub_type_subscription: values[13].as_i64().unwrap(),
                sub_type_adsupported: values[14].as_i64().unwrap(),
                sub_type_midtier: values[15].as_i64().unwrap(),
            };
            final_results.push(item);
        }
    });

    Arc::try_unwrap(final_results)
        .unwrap()
        .into_inner()
        .unwrap()
}

async fn get_items_from_redis_using_pipeline(
    keys: Vec<std::string::String>,
) -> RedisResult<Vec<Vec<u8>>> {
    let mut con = REDIS_CLIENT.get_connection();
    let mut pipe = redis::pipe();
    for key in &keys {
        pipe.cmd("GET").arg(key);
    }
    pipe.query(&mut con)
}

async fn get_items_from_redis(keys: Vec<std::string::String>) -> Result<Vec<Vec<u8>>, RedisError> {
    get_items_from_redis_using_pipeline(keys).await
}

async fn aggregate_results(
    results: &[Item],
    aggregation_type: AggregationType,
    resolution: Resolution,
) -> Vec<ItemTotal> {
    let mut map: FxHashMap<&String, ItemTotal> = FxHashMap::default();

    match aggregation_type {
        AggregationType::Total => {
            for item in results {
                let entry = map
                    .entry(&item.download_activity_date)
                    .or_insert_with(|| ItemTotal {
                        download_activity_date: item.download_activity_date.clone(),
                        streams: 0,
                        streams_with_skips: 0,
                        skips: 0,
                        streams_passive: 0,
                        streams_active: 0,
                        streams_collection: 0,
                        streams_unknown: 0,
                        saves: 0,
                        sub_type_subscription: 0,
                        sub_type_adsupported: 0,
                        sub_type_midtier: 0,
                    });

                entry.streams += item.streams;
                entry.streams_with_skips += item.streams_with_skips;
                entry.skips += item.skips;
                entry.streams_passive += item.streams_passive;
                entry.streams_active += item.streams_active;
                entry.streams_collection += item.streams_collection;
                entry.streams_unknown += item.streams_unknown;
                entry.saves += item.saves;
                entry.sub_type_subscription += item.sub_type_subscription;
                entry.sub_type_adsupported += item.sub_type_adsupported;
                entry.sub_type_midtier += item.sub_type_midtier;
            }
        }
        _ => unimplemented!(),
    }

    // if resolution is low, only return every 5th item
    match resolution {
        Resolution::Low => {
            let mut res: Vec<_> = map
                .into_values()
                .enumerate()
                .filter(|(i, _)| i % 5 == 0)
                .map(|(_, item)| item)
                .collect();
            res.sort_by(|a, b| a.download_activity_date.cmp(&b.download_activity_date));
            res
        }
        _ => {
            let mut res: Vec<ItemTotal> = map.into_values().collect();
            res.sort_by(|a, b| a.download_activity_date.cmp(&b.download_activity_date));
            res
        }
    }
}

async fn get_sound_recording_timeseries(
    params: SoundRecordingTimeSeriesQueryParams,
) -> Result<impl warp::Reply, warp::Rejection> {
    let isrc = params.isrc;
    let start_date = params.start_date;
    let end_date = params.end_date;
    let aggregation_type = params.aggregation_type;
    let resolution = params.resolution;

    let generate_keys_from_isrc_and_start_date_and_end_date_start = Instant::now();
    let keys =
        generate_keys_from_isrc_and_start_date_and_end_date(isrc, start_date, end_date).await;
    info!(
        "time to generate keys from isrc and start_date and end_date: {:?}",
        generate_keys_from_isrc_and_start_date_and_end_date_start.elapsed()
    );
    info!("number of keys to get: {}", keys.len());

    let get_items_from_redis_start = Instant::now();
    let items_from_redis = get_items_from_redis(keys).await;
    info!(
        "time to get items from redis: {:?}",
        get_items_from_redis_start.elapsed()
    );

    let non_agg_results_start = Instant::now();
    let non_agg_results = decode_and_flatten(items_from_redis).await;
    info!(
        "time to decode and flatten items from redis: {:?}",
        non_agg_results_start.elapsed()
    );
    info!(
        "number of items in non-aggregated result: {}",
        non_agg_results.len()
    );

    let aggregation_type: AggregationType = match aggregation_type.as_str() {
        "TOTAL" => AggregationType::Total,
        "BY_PRODUCT" => AggregationType::ByProduct,
        "BY_COUNTRY" => AggregationType::ByCountry,
        "BY_FEED" => AggregationType::ByFeed,
        _ => panic!("Invalid aggregation type"),
    };

    let resolution: Resolution = match resolution.as_str() {
        "LOW" => Resolution::Low,
        "FULL" => Resolution::Full,
        _ => panic!("Invalid resolution"),
    };

    let aggregate_results_start = Instant::now();
    let aggregated_results =
        aggregate_results(&non_agg_results, aggregation_type, resolution).await;
    info!(
        "time to aggregate results: {:?}",
        aggregate_results_start.elapsed()
    );

    Ok(warp::reply::json(&aggregated_results))
}

async fn ingest_to_redis_from_s3() -> Result<impl warp::Reply, warp::Rejection> {
    let region_provider = RegionProviderChain::default_provider().or_else("us-east-1");
    let config = aws_config::from_env().region(region_provider).load().await;
    // let s3_client = Client::new(&config);

    let s3_client = Arc::new(Client::new(&config));
    // list files in the dev-cucumbers bucket, prefix is buvarov/redis_test/

    let list_objects_output = s3_client
        .list_objects_v2()
        .bucket("dev-cucumbers")
        .prefix("buvarov/redis_test/")
        .send()
        .await
        .unwrap();
    let mut keys: Vec<String> = Vec::with_capacity(810);
    for obj in list_objects_output.contents().iter() {
        keys.push(obj.key().unwrap().to_string());
    }
    info!("keys: {:?}", keys);
    info!("number of keys to get: {}", keys.len());

    let mut total_counter = 0;

    let ingestion_start = Instant::now();

    let chunk_size = (keys.len() + 19) / 20; // Divide keys into 20 chunks, rounding up

    let mut handles = Vec::new();

    for chunk in keys.chunks(chunk_size) {
        let chunk = chunk.to_owned();
        let s3_client = Arc::clone(&s3_client);
        let mut con = REDIS_CLIENT.get_connection();
        let mut pipe = redis::pipe();

        let handle = tokio::spawn(async move {
            for key in chunk {
                let object = s3_client
                    .get_object()
                    .bucket("dev-cucumbers")
                    .key(key)
                    .send()
                    .await
                    .ok()
                    .unwrap();
                let body_future = tokio::spawn(async {
                    let mut body = object.body.into_async_read();
                    let mut buffer = Vec::new();
                    body.read_to_end(&mut buffer).await.unwrap();
                    buffer
                });
                let body = body_future.await.unwrap();
                let mut decoder = GzDecoder::new(&body[..]);

                let mut buffer_string = String::with_capacity(1024 * 1024 * 500);
                decoder.read_to_string(&mut buffer_string).unwrap();

                let mut counter = 0;

                for line in buffer_string.lines() {
                    let item: serde_json::Value = serde_json::from_str(&line).unwrap();

                    if let Some(obj) = item.as_object() {
                        for (key, value) in obj {
                            let value_str = value.to_string();
                            let mut encoder = GzEncoder::new(Vec::new(), Compression::best());
                            encoder.write_all(value_str.as_bytes()).unwrap();
                            let compressed_bytes = encoder.finish().unwrap();
                            pipe.cmd("SET").arg(key).arg(compressed_bytes).ignore();
                            counter += 1;

                            if counter >= 1000 {
                                pipe.execute(&mut con);
                                pipe.clear();
                                info!("executed 1000 commands in the pipeline");
                                counter = 0;
                                total_counter += 1000;
                            }
                        }
                    }
                }

                // Execute any remaining commands in the pipeline
                if counter > 0 {
                    pipe.execute(&mut con);
                    info!("executed remaining commands in the pipeline")
                }
            }
        });

        handles.push(handle);
    }

    // Wait for all tasks to complete
    for handle in handles {
        handle.await.unwrap();
    }

    info!("total number of keys ingested: {}", total_counter);

    info!(
        "time to ingest to redis from s3: {:?}",
        ingestion_start.elapsed()
    );

    Ok(warp::reply::json(&"Success"))
}

#[derive(Deserialize)]
struct SoundRecordingTimeSeriesQueryParams {
    isrc: String,
    start_date: String,
    end_date: String,
    aggregation_type: String,
    resolution: String,
}

#[tokio::main]
async fn main() {
    env_logger::init();

    let ip_string = std::env::var("IP").unwrap_or_else(|_| String::from("127.0.0.1"));
    let ip: [u8; 4] = ip_string
        .split('.')
        .map(|s| s.parse().expect("IP must be a valid IPv4 address"))
        .collect::<Vec<_>>()
        .try_into()
        .expect("IP must be a valid IPv4 address");
    let port_string = std::env::var("PORT").unwrap_or_else(|_| String::from("8085"));
    let port: u16 = port_string.parse().expect("PORT must be a number");

    let handler = warp::path("rust-sound-recording-timeseries")
        .and(warp::query::<SoundRecordingTimeSeriesQueryParams>())
        .and_then(get_sound_recording_timeseries)
        .or(warp::path("ingest-to-redis-from-s3").and_then(ingest_to_redis_from_s3));

    warp::serve(handler).run((ip, port)).await;
}