use bstr::ByteSlice;
use crate::common::{diff::Hunk, io::WriteTo, lines::Lines, lines::LinesSlice};
use smol::lock::Mutex;
use std::{mem, sync::LazyLock};
use super::{frontend::Frontend, http};

static RPC_MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| { Mutex::new(()) });

// Define an alias for the callback function signature. Since Rust does not support type
// aliases for traits, we need to workaround that by a blanket implementation of a custom
// trait.
pub(crate) trait Callback: FnOnce(Result<(Vec<Lines>, Vec<u8>), String>) + 'static {}
impl<U> Callback for U where U: FnOnce(Result<(Vec<Lines>, Vec<u8>), String>) + 'static {}

struct Request {
    patch: Vec<u8>,
    code: Vec<u8>,
    callback: Box<dyn Callback>,
    id: usize,
}

#[derive(Default, derive_more::Deref, derive_more::DerefMut)]
struct RequestQueue(Vec<Request>);

pub struct Rpc {
    url: String,
    batch_size: usize,
    last_id: usize,
    queue: RequestQueue,
}

impl Request {
    fn new<F: Callback>(callback: F, id: usize) -> Self {
        Self {
            patch: vec![],
            code: vec![],
            callback: Box::new(callback),
            id,
        }
    }

    fn to_json_value(&self) -> serde_json::Value {
        serde_json::json!({
            "jsonrpc": "2.0",
            "method": "inference",
            "id": self.id,
            "params": {
                "patch": self.patch.to_str_lossy().into_owned(),
                "code": self.code.to_str_lossy().into_owned(),
            },
        })
    }
}

impl RequestQueue {
    /// Dispatch the error message to all requests.
    fn dispatch_error(self, msg: String) {
        for request in self.0 {
            (request.callback)(Err(msg.clone()));
        }
    }

    fn parse_response(data: &serde_json::Value) -> Result<(Vec<Lines>, Vec<u8>), String> {
        if let Some(result) = data.get("result") {
            if let Some(array) = result.as_array() {
                if array.is_empty() {
                    return Err("Invalid format of the returned data.".to_owned());
                }
                let mut lines = vec![];
                let mut confidence = vec![];
                for pred in array {
                    if let Some(pred_array) = pred.as_array() {
                        if pred_array.len() != 2 {
                            return Err("Invalid format of the returned data.".to_owned());
                        }
                        let code = match pred_array[0].as_str() {
                            None => return Err("Invalid format of the returned data.".to_owned()),
                            Some(code) => code
                        };
                        lines.push(code.as_bytes().to_owned().into());
                        let score = match pred_array[1].as_f64() {
                            None => return Err("Invalid format of the returned data.".to_owned()),
                            Some(score) => score
                        };
                        confidence.push((1_000_000_f64.powf(score) * 100.).max(0.).min(100.) as u8);
                    } else {
                        return Err("Invalid format of the returned data.".to_owned());
                    }
                }
                return Ok((lines, confidence));
            } else {
                return Err("Invalid format of the returned data.".to_owned());
            }
        } else if let Some(error) = data.get("error") {
            if let Some(error_obj) = error.as_object() {
                if let Some(message) = error_obj.get("message") {
                    if let Some(message_str) = message.as_str() {
                        return Err(message_str.to_owned());
                    }
                }
            }
        }
        Err("Invalid format of the returned data.".to_owned())
    }

    /// Find the corresponding request and give the response to it. Ignore the response if
    /// there's no corresponding request.
    fn response(&mut self, data: &serde_json::Value) {
        if !data.is_object() {
            return;
        }
        let id = match data.get("id") {
            Some(id_val) => match id_val.as_u64() {
                Some(id_u64) => id_u64 as usize,
                None => return,
            },
            None => return,
        };
        let pos = match self.iter().position(|r| r.id == id) {
            None => return,
            Some(pos) => pos
        };
        let request = self.remove(pos);
        (request.callback)(Self::parse_response(data));
    }
}

impl Rpc {
    pub fn new(url: String, batch_size: usize) -> Self {
        Self {
            url,
            batch_size,
            queue: RequestQueue::default(),
            last_id: 0,
        }
    }

    pub(crate) fn queue_ai<F: Callback>(
        &mut self,
        lines_slice: LinesSlice,
        hunk: &Hunk,
        f: F,
        frontend: impl Frontend + 'static,
    ) {
        let mut request = Request::new(f, self.last_id);
        self.last_id += 1;
        hunk.write_to(&mut request.patch).unwrap();
        lines_slice.write_to(&mut request.code).unwrap();
        self.queue.push(request);
        if self.queue.len() == self.batch_size {
            self.execute(frontend);
        }
    }

    pub(crate) fn execute(&mut self, frontend: impl Frontend + 'static) {
        if self.queue.len() == 0 {
            return;
        }
        let queue = mem::replace(&mut self.queue, RequestQueue::default());
        let url = self.url.clone();
        frontend.clone().spawn_future(async move {
            process_requests(url, queue, frontend).await
        });
    }
}

fn response_to_json_value(data: Result<Vec<u8>, String>)
    -> Result<serde_json::Value, String>
{
    let data = data?;
    let parsed: serde_json::Value = serde_json::from_slice(&data)
        .map_err(|e| e.to_string())?;
    if !parsed.is_array() {
        return Err("Invalid format of the returned data.".to_owned());
    }
    Ok(parsed)
}

async fn process_requests(url: String, mut queue: RequestQueue, frontend: impl Frontend) {
    let mut data = Vec::new();
    for request in queue.iter() {
        data.push(request.to_json_value());
    }
    let data_array = serde_json::Value::Array(data);

    // serialize all RPC calls for backend usage fairness
    let _guard = RPC_MUTEX.lock().await;
    let response_array = http::post(
        &url,
        serde_json::to_string(&data_array).unwrap(),
        &frontend,
    ).await;

    let response_array = match response_to_json_value(response_array) {
        Err(msg) => {
            queue.dispatch_error(msg);
            return;
        }
        Ok(data) => data
    };
    if let Some(responses) = response_array.as_array() {
        for response in responses {
            queue.response(response);
        }
    }
    if queue.len() > 0 {
        queue.dispatch_error("No data returned from the server.".to_owned())
    }
}
