use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
use crate::net::encoding::encode_board_state_perspective;
use crate::net::model::{ChessModel, TrainingSample};
use crate::train::{apply_temperature, sample_move};
use burn::tensor::Device;
use burn::Tensor;
use chess::Color;
use moka::sync::Cache;
use rand::prelude::SmallRng;
use rand::SeedableRng;
use std::collections::VecDeque;
use std::hash::Hash;
use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender};
use std::time::Duration;

pub struct Worker {
    id: usize,
    num_games: usize,
    // games: Vec<BoardState>,
    request_tx: Sender<NNRequest>,
    data_tx: Sender<WorkerData>,
    curr_request_id: u64,
    device: Device,
    mcts_config: MctsConfig,
}

pub struct WorkerData {
    pub(crate) worker_id: usize,
    pub(crate) data: Vec<TrainingSample>,
}

impl Worker {
    pub fn new(
        id: usize,
        num_games: usize,
        sender: Sender<NNRequest>,
        data_tx: Sender<WorkerData>,
        device: Device,
        mcts_config: MctsConfig,
    ) -> Worker {
        Worker {
            id,
            // games: (0..num_games)
            //     .map(|_| BoardState::default())
            //     .collect::<Vec<BoardState>>(),
            num_games,
            request_tx: sender,
            data_tx,
            curr_request_id: 0,
            device,
            mcts_config,
        }
    }

    async fn run_game<F, Fut>(
        worker_id: usize,
        mcts_config: MctsConfig,
        model_inference: F,
        data_tx: Sender<WorkerData>,
    ) where
        F: Fn(&BoardState) -> Fut,
        Fut: Future<Output = (Vec<f32>, f32)>,
    {
        // println!("game started");
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap();
        let seed = now.as_nanos() as u64;
        let mut rng = SmallRng::seed_from_u64(seed);
        let mut mcts: Mcts = Mcts {
            config: mcts_config,
        };

        let mut board_state = BoardState::default();
        let mut episode_buffer: Vec<MctsResults> = vec![];

        let mut game_hist = String::new();

        let mut ply_count: usize = 0;

        while board_state.status == BoardStateStatus::Ongoing {
            // let before = Instant::now();
            let results = mcts.search(&board_state, &model_inference, true).await;
            // println!("{}", before.elapsed().as_millis());
            episode_buffer.push(results);

            let temp = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3);
            // let temp = if ply_count < 30 { 1.0 } else { 0.2 };

            let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);

            let mv = sample_move(&adjusted, &mut rng, &board_state.board).unwrap();

            game_hist.push_str(mv.to_string().as_str());
            game_hist.push(' ');

            board_state.apply_move(mv);
            ply_count += 1;

            // truncate long games
            if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
                board_state.status = BoardStateStatus::FiftyMove;
            }
        }

        println!("game finished: {}", game_hist);

        let mut samples: Vec<TrainingSample> = Vec::new();
        for result in episode_buffer.iter() {
            let stm = result.board_state.board.side_to_move();
            let outcome = match board_state.status {
                BoardStateStatus::WhiteWinner => {
                    if stm == Color::White {
                        1.0
                    } else {
                        -1.0
                    }
                }
                BoardStateStatus::BlackWinner => {
                    if stm == Color::Black {
                        1.0
                    } else {
                        -1.0
                    }
                }
                _ => 0.0,
            };

            let blended_target = 0.75 * outcome + 0.25 * result.value;

            samples.push(TrainingSample::from_mcts_with_outcome(
                result.clone(),
                blended_target,
            ));
        }

        let worker_data = WorkerData {
            worker_id,
            data: samples,
        };
        data_tx.send(worker_data).unwrap();
    }

    // async fn model_inference(&self, board_state: BoardState) -> (Vec<f32>, f32) {
    //     let state = encode_board_state_perspective(&board_state, &self.device).unsqueeze::<4>();
    //     let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
    //
    //     self.request_tx
    //         .send(NNRequest {
    //             request_tensor: state,
    //             response_tx: tx,
    //         })
    //         .unwrap();
    //
    //     let response = rx.await.expect("worker dropped inference response");
    //     let p = response
    //         .response_tensors
    //         .0
    //         .into_data()
    //         .to_vec::<f32>()
    //         .unwrap();
    //     let v = response
    //         .response_tensors
    //         .1
    //         .into_data()
    //         .to_vec::<f32>()
    //         .unwrap()[0];
    //
    //     (p, v)
    // }

    pub async fn run_worker(self, cmd_rx: Receiver<WorkerCmd>) {
        let mut tasks = Vec::new();
        // println!("Worker started");
        for _ in 0..self.num_games {
            let mcts_config = self.mcts_config.clone();
            let data_tx = self.data_tx.clone();
            let request_tx = self.request_tx.clone();
            let device = self.device.clone();

            tasks.push(tokio::spawn(async move {
                Self::run_game(
                    self.id,
                    mcts_config,
                    |board_state| {
                        let request_tx = request_tx.clone();
                        let device = device.clone();

                        model_inference_static(board_state.clone(), request_tx, device)
                    },
                    data_tx,
                )
                .await;
            }));
        }

        loop {
            // Remove completed tasks.
            tasks.retain(|task| !task.is_finished());

            // If every game has completed, we're done.
            if tasks.is_empty() {
                break;
            }

            match cmd_rx.recv_timeout(Duration::from_millis(10)) {
                Ok(cmd) => match cmd {
                    WorkerCmd::Shutdown => break,

                    WorkerCmd::Dispatch => {
                        let mcts_config = self.mcts_config.clone();
                        let request_tx = self.request_tx.clone();
                        let device = self.device.clone();
                        let data_tx = self.data_tx.clone();

                        tasks.push(tokio::spawn(async move {
                            Self::run_game(
                                self.id,
                                mcts_config,
                                |board_state| {
                                    let request_tx = request_tx.clone();
                                    let device = device.clone();

                                    model_inference_static(board_state.clone(), request_tx, device)
                                },
                                data_tx,
                            )
                            .await;
                        }));
                    }
                },

                Err(RecvTimeoutError::Timeout) => {
                    continue;
                }

                Err(RecvTimeoutError::Disconnected) => {
                    println!("Request channel disconnected");
                    break;
                }
            }
        }

        for task in tasks {
            task.await.expect("Worker game task panicked");
        }

        println!("Worker {} shut down", self.id);
    }
}

async fn model_inference_static(
    board_state: BoardState,
    request_tx: Sender<NNRequest>,
    device: Device,
) -> (Vec<f32>, f32) {
    let cache_key = board_state.nn_cache_key();
    let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>();

    let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();

    request_tx
        .send(NNRequest {
            cache_key,
            request_tensor: state,
            response_tx: tx,
        })
        .unwrap();

    let response = rx.await.expect("worker dropped inference response");

    (response.policy, response.value)
}

pub enum NNServerCmd {
    Shutdown,
}

pub enum WorkerCmd {
    Shutdown,
    Dispatch,
}

pub struct NNServer {
    command_rx: Receiver<NNServerCmd>,
    request_rx: Receiver<NNRequest>,
    model: ChessModel,
    device: Device,
    batch_size: usize,
    cache_size: u64,
}

impl NNServer {
    pub fn new(
        command_rx: Receiver<NNServerCmd>,
        request_rx: Receiver<NNRequest>,
        model: ChessModel,
        device: Device,
        batch_size: usize,
        cache_size: u64,
    ) -> NNServer {
        NNServer {
            command_rx,
            request_rx,
            model,
            device,
            batch_size,
            cache_size,
        }
    }
}
#[derive(Debug, Clone, Eq, Hash, PartialEq)]
struct NNJobData {
    worker_id: usize,
    request_id: u64,
    // game_id: u32,
}
pub struct NNRequest {
    cache_key: NNCacheKey,
    request_tensor: Tensor<4>,
    response_tx: tokio::sync::oneshot::Sender<NNResponse>,
}

#[derive(Debug, Clone)]
pub struct NNResponse {
    policy: Vec<f32>,
    value: f32,
}

impl NNServer {
    pub fn run_server(&self) {
        let mut cache: Cache<NNCacheKey, NNResponse> = Cache::new(self.cache_size);
        let mut queue: VecDeque<NNRequest> = VecDeque::new();
        let mut request_counter: usize = 0;
        let mut cache_hits: usize = 0;
        let mut batches_ran: usize = 0;

        loop {
            // Check for shutdown before doing any work.
            if let Ok(cmd) = self.command_rx.try_recv() {
                match cmd {
                    NNServerCmd::Shutdown => {
                        println!(
                            "NN server shutting down. Served {} requests, with {} cache hits ({:2}%). Ran {} batches (avg {}requests/batch).",
                            request_counter,
                            cache_hits,
                            cache_hits as f64 / request_counter as f64 * 100.0,
                            batches_ran,
                            (request_counter - cache_hits) as f64 / batches_ran as f64,
                        );

                        break;
                    } // other commands...
                }
            }

            // If we have no queued requests, wait for either a request
            // or periodically wake up to check for shutdown.
            if queue.is_empty() {
                match self.request_rx.recv_timeout(Duration::from_millis(10)) {
                    Ok(first) => {
                        request_counter += 1;
                        if try_use_cache(first, &mut cache, &mut queue) {
                            cache_hits += 1;
                            continue;
                        }
                    }

                    Err(RecvTimeoutError::Timeout) => {
                        continue;
                    }

                    Err(RecvTimeoutError::Disconnected) => {
                        println!("Request channel disconnected");
                        break;
                    }
                }
            }

            while let Ok(req) = self.request_rx.try_recv() {
                request_counter += 1;
                if try_use_cache(req, &mut cache, &mut queue) {
                    cache_hits += 1;
                    continue;
                }
            }

            let mut jobs_to_send: Vec<NNRequest> = Vec::new();

            for i in 0..self.batch_size {
                if let Some(job) = queue.pop_front() {
                    jobs_to_send.push(job);
                } else {
                    break;
                }
            }

            let batched_tensor = Tensor::cat(
                jobs_to_send
                    .iter()
                    .map(|j| j.request_tensor.clone())
                    .collect(),
                0,
            );

            // println!("Sending {} jobs to gpu", jobs_to_send.len());

            let (p, v) = self.model.forward(batched_tensor);
            batches_ran += 1;
            let p_split = p.chunk(jobs_to_send.len(), 0);
            let v_split = v.chunk(jobs_to_send.len(), 0);

            for (i, job) in jobs_to_send.into_iter().enumerate() {
                let p = p_split[i].clone().into_data().to_vec::<f32>().unwrap();

                let v = v_split[i].clone().into_data().to_vec::<f32>().unwrap()[0];

                let res = NNResponse {
                    policy: p,
                    value: v,
                };

                cache.insert(job.cache_key.clone(), res.clone());

                job.response_tx
                    .send(res)
                    .expect("worker dropped inference request");
            }
        }
    }
}

fn try_use_cache(
    req: NNRequest,
    cache: &mut Cache<NNCacheKey, NNResponse>,
    queue: &mut VecDeque<NNRequest>,
) -> bool {
    if let Some(cached) = cache.get(&req.cache_key) {
        // println!("Cache hit");
        req.response_tx.send(cached).unwrap();
        true
    } else {
        queue.push_back(req);
        false
    }
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub struct NNCacheKey {
    pub board_hash: u64,
    pub halfmove_clock: u8,
    pub repetition_two: bool,
}