use std::{collections::HashSet, fmt, fs};
use directories::ProjectDirs;
use expanduser::expanduser;
use serde::{Deserialize, Serialize};
use walkdir::WalkDir;
use zstd;

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

fn default_state_path() -> String {
    let name = "mood.bin.zstd";
    ProjectDirs::from("qualifier", "organisation", "mood")
        .and_then(|pd| pd.config_dir().join(name).to_str().map(|x| x.to_string()))
        .unwrap_or(format!("~/.config/mood/{name}"))
}

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

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

pub struct UpdateResult {
    added: usize,
    removed: usize,
    new_total: usize,
}

impl fmt::Display for UpdateResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "+{}/-{}, total {}", self.added, self.removed, self.new_total)
    }
}

impl State {
    pub fn try_load() -> Result<State, String> {
        expanduser(default_state_path())
            .map_err(|e| e.to_string())
            .and_then(|p| fs::read(p).map_err(|e| e.to_string()))
            .and_then(|contents| zstd::decode_all(&contents[..]).map_err(|e| e.to_string()))
            .and_then(|decompressed| bincode::deserialize(&decompressed).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(std::fs::create_dir_all);
        let serialised =
            bincode::serialize(self).map_err(|e| format!("Failed to serialize state: {}", e))?;
        let compressed = zstd::encode_all(serialised.as_slice(), 3)
            .map_err(|e| format!("Failed to compress state: {}", e))?;
        fs::write(&path, compressed)
            .map_err(|e| format!("Failed to write state file at {}: {}", str_path, e))?;
        Ok(())
    }

    pub fn update_albums(&mut self, config: &Config) -> UpdateResult {
        let roots: Vec<_> = 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, config) {
                    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(entry: &walkdir::DirEntry, config: &Config) -> 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(), config))
            })
        })
    }

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

impl State {
    pub fn dump_albums(&self) -> Result<String, String> {
        serde_json::to_string_pretty(&self.albums).map_err(|e| e.to_string())
    }

    pub fn dump_learner(&self) -> Result<String, String> {
        serde_json::to_string_pretty(&self.learner).map_err(|e| e.to_string())
    }

    pub fn dump_trajectory(&self) -> Result<String, String> {
        serde_json::to_string_pretty(&self.trajectory).map_err(|e| e.to_string())
    }
}