ck3-regions 0.0.5

use std::{
    collections::HashMap,
    hash::Hash,
    fmt::Display,
};

use grid::Grid;
use clap::ValueEnum;

//
// Exports
//

pub use self::error::{TitleIdError, MapDataError};

//
// Interface aliases
//

pub type Rgb = image::Rgb<u8>;

//
// pub enum TitleTier
//

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, ValueEnum)]
pub enum TitleTier {
    Barony,
    County,
    Duchy,
    Kingdom,
    Empire
}

impl TitleTier {
    pub const TIERS_COUNT: usize = 5;

    pub fn next_tier(self) -> Option<Self> {
        use TitleTier::*;

        match self {
            Barony  => Some(County),
            County  => Some(Duchy),
            Duchy   => Some(Kingdom),
            Kingdom => Some(Empire),
            Empire  => None
        }
    }
}

//
// pub struct TitleId
//

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TitleId {
    id_str: String,
    tier:   TitleTier
}

impl Hash for TitleId {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.id_str.hash(state);
    }
}

impl TryFrom<String> for TitleId {
    type Error = TitleIdError;

    fn try_from(id_str: String) -> Result<Self, Self::Error> {
        if id_str.is_empty() {
            return Err(TitleIdError::EmptyIdString);
        }

        if let Some((prefix, body)) = id_str.split_once('_') {
            if let Some(tier) = match prefix {
                "b" => Some(TitleTier::Barony),
                "c" => Some(TitleTier::County),
                "d" => Some(TitleTier::Duchy),
                "k" => Some(TitleTier::Kingdom),
                "e" => Some(TitleTier::Empire),
                _ => None
            } {
                if body.is_empty() {
                    Err(Self::Error::MissingIdBody(id_str))
                } else {
                    Ok(Self{id_str, tier})
                }
            } else {
                Err(Self::Error::InvalidTierPrefix{id_str: id_str.clone(), prefix_range: ..prefix.len()})
            }
        } else {
            Err(Self::Error::MissingTierPrefix(id_str))
        }
    }
}

impl TryFrom<&str> for TitleId {
    type Error = TitleIdError;

    fn try_from(id_str: &str) -> Result<Self, Self::Error> {
        String::from(id_str).try_into()
    }
}

impl AsRef<str> for TitleId {
    fn as_ref(&self) -> &str {
        self.id_str.as_ref()
    }
}

impl Display for TitleId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.id_str)
    }
}

impl TitleId {
    pub fn get_tier(&self) -> TitleTier {
        self.tier
    }
}

//
// pub struct ProvinceId
//

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ProvinceId(pub u32);

impl ProvinceId {
    pub const NULL_ID: ProvinceId = ProvinceId(0);
}

//
// pub struct MapData
//

pub struct MapData {
    containing_title_lookup_table: HashMap<(TitleId, TitleTier), TitleId>,
    province_baronies:   HashMap<ProvinceId, TitleId>,
    province_id_by_rgb:    HashMap<Rgb, ProvinceId>
}

impl MapData {
    //
    // Interface
    //

    pub fn get_province_id_by_rgb(&self, rgb: Rgb) -> Option<ProvinceId> {
        self.province_id_by_rgb.get(&rgb).copied()
    }

    pub fn get_containing_title_id_by_province_id(&self, province_id: ProvinceId, target_tier: TitleTier) -> Option<&TitleId> {
        let barony_id = self.get_barony_by_province_id(province_id)?;

        self.get_containing_title_id(barony_id, target_tier)
    }

    //
    // Service
    //

    fn get_barony_by_province_id(&self, province_id: ProvinceId) -> Option<&TitleId> {
        self.province_baronies.get(&province_id)
    }

    fn get_containing_title_id(&self, base_title_id: &TitleId, target_tier: TitleTier) -> Option<&TitleId> {
        self.containing_title_lookup_table.get(&(base_title_id.clone(), target_tier))
    }
}

//
// pub struct MapDataBuilder
//

pub struct MapDataBuilder {
    rgb_by_province_id: HashMap<ProvinceId, Rgb>,
    province_baronies:  HashMap<ProvinceId, TitleId>,
    liege_by_title_id:  HashMap<TitleId, Option<TitleId>>,
    errors:             Vec<MapDataError>,
}

impl MapDataBuilder {
    pub fn new() -> Self {
        Self{
            rgb_by_province_id: HashMap::new(),
            province_baronies:  HashMap::new(),
            liege_by_title_id:  HashMap::new(),
            errors:             Vec::new(),
        }
    }

    pub fn add_province(&mut self, id: ProvinceId, rgb: Rgb) -> &mut MapDataBuilder {
        if let Some(_) = self.rgb_by_province_id.get(&id) {
            self.errors.push(MapDataError::DuplicateProvinceId(id));
        } else if let Some((&old_id, _)) = self.rgb_by_province_id.iter().find(|(&_, &old_rgb)| old_rgb == rgb) {
            self.errors.push(MapDataError::DuplicateProvinceRgb{rgb, old_id, new_id: id})
        } else {
            self.rgb_by_province_id.insert(id, rgb);
        }

        self
    }

    pub fn add_title(
        &mut self,
        id:          TitleId,
        liege_id:    Option<TitleId>,
        province_id: Option<ProvinceId>
    ) -> &mut MapDataBuilder {
        if self.liege_by_title_id.contains_key(&id) {
            self.errors.push(MapDataError::DuplicateTitleId(id));

            return self;
        }

        if let Some(province_id) = province_id {
            if id.get_tier() != TitleTier::Barony {
                self.errors.push(MapDataError::NonBaronyHasProvince{province_id, title_id: id});

                return self;
            }

            if let Some(old_barony_id) = self.province_baronies.get(&province_id) {
                self.errors.push(MapDataError::AmbiguousProvinceBarony{province_id, old_barony_id: old_barony_id.clone(), new_barony_id: id});

                return self;
            }
        } else if id.get_tier() == TitleTier::Barony {
            self.errors.push(MapDataError::MissingBaronyProvince(id));

            return self;
        }

        if liege_id.as_ref().map_or(false, |liege_id| liege_id.get_tier() <= id.get_tier()) {
            self.errors.push(MapDataError::InvalidLiegeTier{id, liege_id: liege_id.expect("confirmed Some by map_or()")});

            return self;
        }

        province_id.map(|province_id| self.province_baronies.insert(province_id, id.clone()));
        self.liege_by_title_id.insert(id, liege_id);

        self
    }

    pub fn build(mut self) -> Result<MapData, Vec<MapDataError>> {
        let mut containing_title_lookup_table = HashMap::new();

        for (title_id, liege_id) in &self.liege_by_title_id {
            // Each title is its own containing title of its own tier
            containing_title_lookup_table.insert((title_id.clone(), title_id.get_tier()), title_id.clone());

            // Find higher tier containing titles
            let mut current_title_id = title_id;
            loop {
                let current_liege_id = if current_title_id == title_id {
                    liege_id
                } else {
                    if let Some(liege_id) = self.liege_by_title_id.get(&current_title_id) {
                        liege_id
                    } else {
                        self.errors.push(MapDataError::UnknownLiegeTitleId(current_title_id.clone()));

                        &None
                    }
                };

                if let Some(current_liege_id) = current_liege_id {
                    assert!(current_liege_id.get_tier() > current_title_id.get_tier());

                    containing_title_lookup_table.insert(
                        (title_id.clone(), current_liege_id.get_tier()), current_liege_id.clone()
                    );

                    current_title_id = current_liege_id;
                } else {
                    break;
                }
            }
        }

        self.province_baronies.keys()
            .filter(|province_id| !self.rgb_by_province_id.contains_key(&province_id))
            .for_each(|&province_id| self.errors.push(MapDataError::MissingProvinceRgb(province_id)));

        if !self.errors.is_empty() {
            return Err(self.errors);
        }

        let province_id_by_rgb = self.rgb_by_province_id
            .drain()
            .map(|(province_id, rgb)| (rgb, province_id))
            .collect();

        Ok(MapData{
            containing_title_lookup_table,
            province_id_by_rgb,
            province_baronies: self.province_baronies,
        })
    }
}

//
// pub struct ProvincesMap
//

pub struct ProvincesMap(pub Grid<ProvinceId>);

//
// Errors
//

mod error {
    use std::ops::RangeTo;

    use thiserror::Error;

    use super::*;

    //
    // pub enum TitleIdError
    //

    #[derive(Error, Debug)]
    pub enum TitleIdError {
        #[error("empty title ID string")]
        EmptyIdString,
        #[error("title ID {0:?} is missing the tier prefix; tier prefix must be separated by '_'")]
        MissingTierPrefix(String),
        #[error("title ID {id_str:?} has unrecognized tier prefix {:?}", &id_str[..prefix_range.end])]
        InvalidTierPrefix{
            id_str: String,
            prefix_range: RangeTo<usize>
        },
        #[error("title ID {0:?} contains nothing but the tier prefix")]
        MissingIdBody(String)
    }

    //
    // pub enum MapDataError
    //

    #[derive(Error, Debug)]
    pub enum MapDataError {
        #[error("duplicate province ID {}", .0.0)]
        DuplicateProvinceId(ProvinceId),
        #[error("duplicate RGB {:?} for province ID {new_id:?}, previously seen with ID {old_id:?}", .rgb.0)]
        DuplicateProvinceRgb{
            rgb:    Rgb,
            new_id: ProvinceId,
            old_id: ProvinceId
        },
        #[error("duplicate title ID {0}")]
        DuplicateTitleId(TitleId),
        #[error(
            "{id} ({:?}) must have a lower tier than its liege {liege_id} ({:?})", id.get_tier(), liege_id.get_tier()
        )]
        InvalidLiegeTier{
            id:       TitleId,
            liege_id: TitleId
        },
        #[error("{:?} title {title_id} has province ID {province_id:?} specified; only baronies can have a province", title_id.get_tier())]
        NonBaronyHasProvince{
            title_id:    TitleId,
            province_id: ProvinceId
        },
        #[error("barony title {0} has no province specified")]
        MissingBaronyProvince(TitleId),
        #[error("province {} specified for {new_barony_id} is already used for {old_barony_id}", .province_id.0)]
        AmbiguousProvinceBarony{
            province_id:   ProvinceId,
            old_barony_id: TitleId,
            new_barony_id: TitleId
        },
        #[error("unknown title {0} specified as liege for another title")]
        UnknownLiegeTitleId(TitleId),
        #[error("province ID {0:?} is used but has no RGB defined")]
        MissingProvinceRgb(ProvinceId)
    }
}

//
// Unit tests
//

#[cfg(test)]
mod tests {
    use thiserror::Error;

    use super::*;

    #[test]
    fn valid_map_data_build_is_ok() {
        assert!(make_valid_test_map_data().is_ok());
    }

    #[test]
    fn province_id_by_rgb_known() {
        let map_data = make_valid_test_map_data().unwrap();

        assert_eq!(map_data.get_province_id_by_rgb(Rgb::from([1, 0, 0])), Some(ProvinceId(1)));
        assert_eq!(map_data.get_province_id_by_rgb(Rgb::from([55, 23, 19])), Some(ProvinceId(3)));
    }

    #[test]
    fn province_id_by_rgb_unknown() {
        let map_data = make_valid_test_map_data().unwrap();

        assert_eq!(map_data.get_province_id_by_rgb(Rgb::from([15, 34, 9])), None);
        assert_eq!(map_data.get_province_id_by_rgb(Rgb::from([99, 99, 99])), None);
    }

    #[test]
    fn containing_title_same_tier() {
        let map_data = make_valid_test_map_data().unwrap();

        let barony_id  = "b_barony_0".try_into().unwrap();
        let county_id  = "c_county_0".try_into().unwrap();
        let duchy_id   = "d_duchy_0".try_into().unwrap();
        let kingdom_id = "k_kingdom_0".try_into().unwrap();
        let empire_id  = "e_empire_0".try_into().unwrap();

        let containing_barony_id  = map_data.get_containing_title_id(&barony_id, TitleTier::Barony);
        let containing_county_id  = map_data.get_containing_title_id(&county_id, TitleTier::County);
        let containing_duchy_id   = map_data.get_containing_title_id(&duchy_id, TitleTier::Duchy);
        let containing_kingdom_id = map_data.get_containing_title_id(&kingdom_id, TitleTier::Kingdom);
        let containing_empire_id  = map_data.get_containing_title_id(&empire_id, TitleTier::Empire);

        assert_eq!(containing_barony_id, Some(&barony_id));
        assert_eq!(containing_county_id, Some(&county_id));
        assert_eq!(containing_duchy_id, Some(&duchy_id));
        assert_eq!(containing_kingdom_id, Some(&kingdom_id));
        assert_eq!(containing_empire_id, Some(&empire_id));
    }

    #[test]
    fn containing_title_higher_tier() {
        let map_data = make_valid_test_map_data().unwrap();

        assert_eq!(
            map_data.get_containing_title_id(&"b_barony_0".try_into().unwrap(), TitleTier::County),
            Some(&"c_county_0".try_into().unwrap())
        );
        assert_eq!(
            map_data.get_containing_title_id(&"b_barony_0".try_into().unwrap(), TitleTier::Duchy),
            Some(&"d_duchy_0".try_into().unwrap())
        );
        assert_eq!(
            map_data.get_containing_title_id(&"b_barony_0".try_into().unwrap(), TitleTier::Kingdom),
            Some(&"k_kingdom_0".try_into().unwrap())
        );
        assert_eq!(
            map_data.get_containing_title_id(&"b_barony_0".try_into().unwrap(), TitleTier::Empire),
            Some(&"e_empire_0".try_into().unwrap())
        );
    }

    //
    // Test types
    //

    #[derive(Error, Debug)]
    enum TestError {
        #[error("{0:?}")]
        InvalidMapData(Vec<MapDataError>),
        #[error("{0}")]
        InvalidTitleId(#[from] TitleIdError)
    }

    impl From<Vec<MapDataError>> for TestError {
        fn from(value: Vec<MapDataError>) -> Self {
            Self::InvalidMapData(value)
        }
    }

    //
    // Test service
    //

    fn make_valid_test_map_data() -> Result<MapData, TestError> {
        let mut map_data_builder = MapDataBuilder::new();
        map_data_builder
            .add_province(ProvinceId(1), Rgb::from([1, 0, 0]))
            .add_province(ProvinceId(2), Rgb::from([42, 105, 18]))
            .add_province(ProvinceId(3), Rgb::from([55, 23, 19]))
            .add_province(ProvinceId(50), Rgb::from([9, 255, 1]))
            .add_title("b_barony_0".try_into()?, Some("c_county_0".try_into()?), Some(ProvinceId(1)))
            .add_title("c_county_0".try_into()?, Some("d_duchy_0".try_into()?), None)
            .add_title("d_duchy_0".try_into()?, Some("k_kingdom_0".try_into()?), None)
            .add_title("k_kingdom_0".try_into()?, Some("e_empire_0".try_into()?), None)
            .add_title("e_empire_0".try_into()?, None, None)
            .add_title("b_barony_1".try_into()?, Some("c_county_0".try_into()?), Some(ProvinceId(2)))
            .add_title("b_barony_2".try_into()?, Some("c_county_1".try_into()?), Some(ProvinceId(3)))
            .add_title("c_county_1".try_into()?, None, None);

        map_data_builder.build().map_err(TestError::from)
    }
}