use std::{collections::HashSet, fs};

use derive_more::Display;
use expanduser::expanduser;
use serde::{Deserialize, Serialize};
use walkdir::WalkDir;

use crate::{
    config::{Config, default_state_path},
    learner::Learner,
    trajectory::Trajectory,
};

#[derive(Serialize, Deserialize, Default)]
pub struct State {
    pub config: Config,
    trajectory: Trajectory,
    learner: Learner,
    albums: HashSet<String>,
}

impl State {
    pub fn next(&mut self) -> Option<(String, f32)> {
        let learning = self.trajectory.step(self.config.skip_window_secs);
        if let Some(ref learning) = learning {
            println!("Obtained learning: {learning:?}");
            self.learner.learn(&self.config, learning);
        }
        if let Some((new_album, prob)) = self.learner.sample(
            self.trajectory.slice(),
            &learning.as_ref().into(),
            &self.albums,
            self.config.temperature,
        ) {
            self.trajectory.log(&new_album);
            Some((new_album, prob))
        } else {
            None
        }
    }
}

#[derive(Display)]
#[display("+{added}/-{removed}, total {new_total}")]
pub struct UpdateResult {
    added: usize,
    removed: usize,
    new_total: usize,
}

impl State {
    pub fn load(path: &str) -> Result<State, String> {
        expanduser(path)
            .map_err(|e| e.to_string())
            .and_then(|p| fs::read_to_string(p).map_err(|e| e.to_string()))
            .and_then(|contents| ron::from_str(&contents).map_err(|e| e.to_string()))
    }

    pub fn try_save(&self) -> Result<(), String> {
        let str_path = default_state_path();
        let path = expanduser(&str_path).map_err(|e| e.to_string())?;
        _ = path.parent().map(|parent| std::fs::create_dir_all(parent));
        let serialised = ron::ser::to_string_pretty(self, ron::ser::PrettyConfig::default())
            .map_err(|e| format!("Failed to serialize state: {}", e))?;
        fs::write(&path, serialised)
            .map_err(|e| format!("Failed to write state file at {}: {}", str_path, e))?;
        Ok(())
    }

    pub fn update_albums(&mut self) -> UpdateResult {
        let roots: Vec<_> = self
            .config
            .music_roots
            .iter()
            .filter_map(|x| expanduser(x).ok().filter(|p| p.is_dir()))
            .collect();

        let mut new_albums: HashSet<String> = HashSet::new();

        for root_path in roots {
            for entry in WalkDir::new(&root_path) {
                let entry = match entry {
                    Ok(e) if e.path().is_dir() => e,
                    _ => continue,
                };

                if !self.contains_audio(&entry) {
                    continue;
                }

                if let Some(rel_str) = entry
                    .path()
                    .strip_prefix(&root_path)
                    .ok()
                    .and_then(|p| p.to_str())
                {
                    new_albums.insert(rel_str.to_string());
                }
            }
        }

        let added = new_albums.difference(&self.albums).count();
        let removed = self.albums.difference(&new_albums).count();
        self.albums.clone_from(&new_albums);

        self.trajectory.prune(&self.albums);
        self.learner.prune(&self.albums);

        UpdateResult {
            added,
            removed,
            new_total: self.albums.len(),
        }
    }

    fn contains_audio(&self, entry: &walkdir::DirEntry) -> bool {
        fs::read_dir(entry.path()).ok().is_some_and(|mut entries| {
            entries.any(|e| {
                e.as_ref()
                    .is_ok_and(|entry| self.is_audio_file(&entry.path()))
            })
        })
    }

    fn is_audio_file(&self, path: &std::path::Path) -> bool {
        !path.is_dir()
            && path
                .extension()
                .and_then(|ext| ext.to_str())
                .map(|s| s.to_lowercase())
                .is_some_and(|ext| self.config.audio_exts.contains(&format!(".{ext}")))
    }
}