use crate::download::helpers::DQTimeseriesRequestArgs;
use crate::download::oauth_client::OAuthClient;
use futures::future;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use std::error::Error;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::runtime::Builder;
use tokio::sync::Semaphore;
use tokio::task;

const API_BASE_URL: &str = "https://api-developer.jpmorgan.com/research/dataquery-authe/api/v2";
const TIMESERIES_ENDPOINT: &str = "/expressions/time-series";

const DQ_THREAD_EXPIRY_SECONDS: u64 = 300;
const DQ_THREAD_SLEEP_MILLIS: u64 = 200;
const DQ_MAX_THREADS: usize = 250;

#[derive(Clone)]
pub struct ParallelRequester {
    oauth_client: Arc<Mutex<OAuthClient>>,
    max_threads: usize,
}

/// Internal implementation of a struct for making parallel requests to the DataQuery API.
impl ParallelRequester {
    pub fn new(oauth_client: OAuthClient) -> Self {
        ParallelRequester {
            oauth_client: Arc::new(Mutex::new(oauth_client)),
            max_threads: DQ_MAX_THREADS,
        }
    }

    /// Builds the headers for the request to the DataQuery API.
    fn build_headers(&self) -> Result<HeaderMap, Box<dyn Error>> {
        let headers_map = self.oauth_client.lock().unwrap().get_headers()?;
        let mut headers = HeaderMap::new();
        for (key, value) in headers_map {
            let header_name = HeaderName::from_bytes(key.as_bytes())?;
            let header_value = HeaderValue::from_str(&value)?;
            headers.insert(header_name, header_value);
        }
        Ok(headers)
    }

    /// Makes an asynchronous request to the DataQuery API.
    async fn _request_async(&self, endpoint: &str) -> Result<String, Box<dyn Error>> {
        let headers = self.build_headers()?;
        let url = format!("{}{}", API_BASE_URL, endpoint);

        let response = reqwest::Client::new()
            .get(&url)
            .headers(headers)
            .send()
            .await?;

        if response.status().is_success() {
            let text = response.text().await?;
            Ok(text)
        } else {
            Err(Box::new(response.error_for_status().unwrap_err()))
        }
    }

    fn _request_expressions(
        &mut self,
        args: DQTimeseriesRequestArgs,
        max_retries: u32,
    ) -> Result<Vec<String>, Box<dyn Error>> {
        let expression_batches: Vec<Vec<String>> = args
            .expressions
            .chunks(20)
            .map(|chunk| {
                let mut vec = chunk.to_vec();
                vec.sort();
                vec
            })
            .collect();

        let prog_batches = match expression_batches.len() {
            0..=250 => 25,
            _ => 100,
        };

        let okay_response_texts = Arc::new(Mutex::new(Vec::new()));
        let failed_batches = Arc::new(Mutex::new(Vec::new()));
        let self_arc = Arc::new(self.clone());

        let runtime = Builder::new_current_thread().enable_all().build()?;
        let semaphore = Arc::new(Semaphore::new(self.max_threads));

        runtime.block_on(async {
            let mut tasks = vec![];
            let mut curr_batch = 1;
            let total_batches = expression_batches.len();
            for batch in expression_batches.into_iter() {
                let self_arc = Arc::clone(&self_arc);
                let okay_response_texts = Arc::clone(&okay_response_texts);
                let failed_batches = Arc::clone(&failed_batches);
                let mut args_clone = args.clone();
                args_clone.update_expressions(batch.clone());
                let ep = format!("{}?{}", TIMESERIES_ENDPOINT, args_clone.as_query_string());
                let permit = semaphore.clone().acquire_owned().await.unwrap();

                if curr_batch % prog_batches == 0 {
                    println!("Requesting batch {} of {}", curr_batch, total_batches);
                }

                let task = task::spawn(async move {
                    let _permit = permit; // Keep the permit alive until the end of the task
                    let res_future = self_arc._request_async(&ep);

                    let res = tokio::time::timeout(
                        Duration::from_secs(DQ_THREAD_EXPIRY_SECONDS),
                        res_future,
                    )
                    .await;
                    match res {
                        Ok(Ok(text)) => {
                            okay_response_texts.lock().unwrap().push(text);
                        }
                        _ => {
                            failed_batches.lock().unwrap().push(batch);
                        }
                    }
                });

                tasks.push(task);

                // Delay before starting the next task
                tokio::time::sleep(Duration::from_millis(DQ_THREAD_SLEEP_MILLIS)).await;
                curr_batch += 1;
            }

            // Await all tasks, ensuring forced termination for any remaining tasks
            let _ = tokio::time::timeout(
                Duration::from_secs(DQ_THREAD_EXPIRY_SECONDS),
                future::join_all(tasks),
            )
            .await;
        });

        // Retry failed batches if any, respecting the max_retries limit
        let failed_batches = failed_batches.lock().unwrap().clone();
        if !failed_batches.is_empty() && max_retries > 0 {
            let mut new_args = args.clone();
            let flattened_failed_batches: Vec<String> =
                failed_batches.into_iter().flatten().collect();
            new_args.update_expressions(flattened_failed_batches);

            let retried_responses = self._request_expressions(new_args, max_retries - 1)?;
            okay_response_texts
                .lock()
                .unwrap()
                .extend(retried_responses);
        } else if !failed_batches.is_empty() && max_retries == 0 {
            return Err("Max retries reached".into());
        }

        // Collect and return successful responses
        let final_responses = okay_response_texts.lock().unwrap().clone();
        Ok(final_responses)
    }

    /// Makes parallel requests to the DataQuery API for the given expressions.
    /// The returned responses are collected and returned as a vector of strings (response texts).
    pub fn request_expressions(
        &mut self,
        args: DQTimeseriesRequestArgs,
        max_retries: u32,
    ) -> Result<Vec<String>, Box<dyn Error>> {
        self._request_expressions(args, max_retries)
    }
}
