use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
use crate::net::encoding::decode_move;
use crate::net::model::{
    ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
};
use burn::data::dataloader::batcher::Batcher;
use burn::module::{AutodiffModule, Module};
use burn::optim::{AdamConfig, GradientsParams, Optimizer};
// use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::prelude::Device;
use burn::store::ModuleRecord;
use chess::{ChessMove, Color};
use rand::rngs::SmallRng;
use rand::seq::SliceRandom;
use rand::{RngExt, SeedableRng};
use std::collections::VecDeque;
use std::fs::File;
use std::io::BufReader;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;
use std::time::{SystemTime, UNIX_EPOCH};

pub struct TrainingConfig {
    pub max_time_s: Option<u64>,
    pub num_iters: Option<usize>,
    pub max_depth: u16, // unused
    pub model_name: String,
    pub load_model: bool,
    pub hidden_channels: usize,
    pub num_blocks: usize,
    pub batch_size: usize,      // num positions sampled to a batch (2048)
    pub num_episodes: usize,    // num games generated per iteration (5000)
    pub buffer_max_size: usize, // max number of samples in the buffer
    pub mcts_config: MctsConfig,
    pub optimizer: AdamConfig,
    pub lr: f64,
    pub seed: Option<u64>,
}

pub fn train(training_config: TrainingConfig, device: &Device) {
    let model_path = format!("artifacts/{}", training_config.model_name.as_str());
    let metadata_path = format!("{}.json", model_path);
    println!("Creating model...");
    let mut model: ChessModel = ChessModelConfig::init(
        training_config.num_blocks,
        training_config.hidden_channels,
        device,
    );
    if training_config.load_model {
        println!("Loading model {}...", model_path);
        // Load model in full precision from MessagePack file
        let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
        model = model.load_record(record)
    }

    let train = Arc::new(AtomicBool::new(true));
    let train_signal = Arc::clone(&train);

    let mut iter: usize = 0;

    if training_config.load_model {
        let file = File::open(&metadata_path).unwrap();

        // 2. Use a buffered reader for efficiency
        let reader = BufReader::new(file);

        // 3. Deserialize JSON directly from the file reader
        let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
        iter = metadata.iterations;
        println!("Loaded model had {} iters", iter);
    }

    let start_time = Instant::now();

    ctrlc::set_handler(move || {
        train_signal.store(false, Ordering::Relaxed);
        println!("Finishing batch before exiting...");
    })
    .expect("Error setting Ctrl-C handler");

    let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();

    let mut mcts: Mcts = Mcts {
        config: training_config.mcts_config,
    };

    // Create RNG once and reuse it for sampling and shuffling
    // Seed from system time (platform default entropy may be unavailable in some contexts)
    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
    let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
    println!("using seed: {}", seed);
    let mut rng = SmallRng::seed_from_u64(seed);

    // Initialize optimizer once so state (moments) persist across steps
    let mut optim = training_config.optimizer.init();

    println!("Starting training...");
    while train.load(Ordering::Relaxed) {
        println!("Iteration: {}", iter);
        let infer_model = model.valid();
        // Gen samples
        println!("Generating {} games...", training_config.num_episodes);
        for episode in 0..training_config.num_episodes {
            println!("Episode: {}", episode);
            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, &infer_model, &device);
                // println!("{}", before.elapsed().as_millis());
                episode_buffer.push(results);

                let temp = if ply_count < 30 {
                    1.0
                } else {
                    0.0
                };

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

                let stm = board_state.board.side_to_move();
                let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
                // println!("playing move: {}", mv);
                if episode == 0 {
                    game_hist.push_str(mv.to_string().as_str());
                    game_hist.push(' ');
                }

                board_state.apply_move(mv);
                ply_count += 1;
                if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
                    board_state.status = BoardStateStatus::FiftyMove;
                }
            }
            if episode == 0 {
                println!("Game history of first game of iteration: {}", game_hist);
            }

            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,
                };

                replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
                    result.clone(),
                    outcome,
                ));

                if replay_buffer.len() > training_config.buffer_max_size {
                    replay_buffer.pop_front();
                }
            }
        }
        // train

        if replay_buffer.is_empty() {
            println!("Replay buffer is empty, skipping training step.");
            continue;
        }

        let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len());
        println!(
            "Finished self-play! Training on {} samples (buffer size: {})...",
            num_samples,
            replay_buffer.len()
        );
        let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
        indices.shuffle(&mut rng);

        let samples = indices
            .into_iter()
            .take(num_samples)
            .map(|i| replay_buffer[i].clone())
            .collect::<Vec<TrainingSample>>();

        let batcher = ChessBatcher {};

        let batch = batcher.batch(samples, &device);

        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);

        let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
        let p_loss_val = output.policy_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
        let v_loss_val = output.value_loss.clone().into_data().to_vec::<f32>().unwrap()[0];

        println!(
            "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
            iter, loss_val, p_loss_val, v_loss_val
        );

        let grads = output.loss.backward();
        let grads = GradientsParams::from_grads(grads, &model);

        model = optim.step(training_config.lr, model, grads);

        iter += 1;

        if iter % 100 == 0 {
            println!("Completed {} iterations", iter);
        }

        if training_config.max_time_s.is_some()
            && start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
        {
            println!("Training stopping due to time limit...");
            break;
        }

        if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
            println!(
                "Training stopping due to iteration limit ({} iters completed)...",
                iter
            );
            break;
        }
    }

    println!("Saving model...");

    let metadata = ModelMetadata {
        name: training_config.model_name,
        iterations: iter,
    };

    std::fs::write(
        metadata_path,
        serde_json::to_string_pretty(&metadata)
            .expect("Should be able to convert metadata to JSON string"),
    )
    .expect("Should be able to write metadata");

    model
        .into_record()
        .save(&model_path)
        .expect("Should be able to save the model");
    println!("Model saved in {:?}, exiting training.", model_path);
    return;
}

fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
    if visits.is_empty() {
        return Vec::new();
    }

    // Special case: deterministic selection
    if temperature == 0.0 {
        let (&best_idx, _) = visits
            .iter()
            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
            .map(|(i, p)| (i, p))
            .unwrap();

        return vec![(best_idx, 1.0)];
    }

    let inv_temp = 1.0 / temperature;

    // Step 1: apply exponent
    let mut adjusted: Vec<(usize, f32)> =
        visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();

    // Step 2: normalize
    let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();

    if sum <= 0.0 {
        return adjusted; // fallback (shouldn't happen in normal MCTS)
    }

    for (_, v) in adjusted.iter_mut() {
        *v /= sum;
    }

    adjusted
}

fn sample_move(
    dist: &[(usize, f32)],
    rng: &mut SmallRng,
    side_to_move: Color,
) -> Option<ChessMove> {
    let mut r: f32 = rng.random_range(0.0..1.0);

    for (idx, p) in dist {
        r -= *p;
        if r <= 0.0 {
            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
        }
    }

    // fallback due to floating point drift
    dist.get(0)
        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
}