factorio_blueprint/factorio-blueprint/src/abstraction.rs

use factorio_core::{
    beltoptions::Beltspeed,
    pathfield::PathField,
    prelude::{Direction, Position, PositionType, Transformable, Transformation},
};
use std::{collections::HashMap, sync::atomic::AtomicUsize};

use crate::{BlueprintEntity, BlueprintPosition};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum UndergroundType {
    Input,
    Output,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum InserterType {
    Burner,
    Normal,
    Fast,
    Bulk,
    Stack,
}

impl InserterType {
    fn string(&self) -> String {
        match self {
            InserterType::Burner => "burner-inserter",
            InserterType::Normal => "inserter",
            InserterType::Fast => "fast-inserter",
            InserterType::Bulk => "bulk-inserter",
            InserterType::Stack => "stack-inserter",
        }
        .to_owned()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ElectricPoleType {
    Small,
    Medium,
    Big,
    Substation,
}

impl ElectricPoleType {
    fn string(&self) -> String {
        match self {
            ElectricPoleType::Small => "small-electric-pole",
            ElectricPoleType::Medium => "medium-electric-pole",
            ElectricPoleType::Big => "big-electric-pole",
            ElectricPoleType::Substation => "substation",
        }
        .to_owned()
    }

    fn size(&self) -> Position {
        match self {
            ElectricPoleType::Small | ElectricPoleType::Medium => Position::new(1, 1),
            ElectricPoleType::Big | ElectricPoleType::Substation => Position::new(2, 2),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Quality {
    Normal,
    Uncommon,
    Rare,
    Epic,
    Legendary,
}

pub struct Entity {
    entity: EntityType,
    position: Position,
    direction: Direction,
    quality: Quality,
}

pub enum EntityType {
    Belt(Beltspeed),
    UndergroundBelt(Beltspeed, UndergroundType),
    Splitter(Beltspeed),
    ElectricPole(ElectricPoleType),
    Inserter {
        inserter_type: InserterType,
        override_stack_size: Option<u8>,
    },
    Unknown {
        name: String,
        size: Position,
        misc: HashMap<String, serde_json::Value>,
    },
}

impl Entity {
    pub fn new(entity: EntityType, position: Position, direction: Direction) -> Self {
        Self {
            entity,
            position,
            direction,
            quality: Quality::Normal,
        }
    }

    pub fn new_belt(beltspeed: Beltspeed, position: Position, direction: Direction) -> Self {
        Self::new(EntityType::Belt(beltspeed), position, direction)
    }

    pub fn new_underground_belt(
        beltspeed: Beltspeed,
        underground_type: UndergroundType,
        position: Position,
        direction: Direction,
    ) -> Self {
        Self::new(
            EntityType::UndergroundBelt(beltspeed, underground_type),
            position,
            direction,
        )
    }

    pub fn new_splitter(beltspeed: Beltspeed, position: Position, direction: Direction) -> Self {
        Self::new(EntityType::Splitter(beltspeed), position, direction)
    }

    pub fn new_inserter(
        inserter_type: InserterType,
        override_stack_size: Option<u8>,
        position: Position,
        direction: Direction,
    ) -> Self {
        Self::new(
            EntityType::Inserter {
                inserter_type,
                override_stack_size,
            },
            position,
            direction,
        )
    }

    pub fn new_electric_pole(electric_pole_type: ElectricPoleType, position: Position) -> Self {
        Self::new(
            EntityType::ElectricPole(electric_pole_type),
            position,
            Direction::Up,
        )
    }

    pub fn new_unknown(
        name: impl AsRef<str>,
        position: Position,
        direction: Direction,
        size: Position,
    ) -> Self {
        Self::new(
            EntityType::Unknown {
                name: name.as_ref().to_owned(),
                size,
                misc: HashMap::new(),
            },
            position,
            direction,
        )
    }

    pub fn quality(mut self, quality: Quality) -> Self {
        self.quality = quality;
        self
    }

    pub fn get_name(&self) -> String {
        match &self.entity {
            EntityType::Belt(beltspeed) => beltspeed.string(),
            EntityType::UndergroundBelt(beltspeed, _underground_type) => {
                beltspeed.string_underground()
            }
            EntityType::Splitter(beltspeed) => beltspeed.string_splitter(),
            EntityType::Unknown {
                name,
                size: _,
                misc: _,
            } => name.clone(),
            EntityType::ElectricPole(electric_pole_type) => electric_pole_type.string(),
            EntityType::Inserter {
                inserter_type,
                override_stack_size: _,
            } => inserter_type.string(),
        }
    }

    pub fn get_maybe_underground_type_string(&self) -> Option<String> {
        match &self.entity {
            EntityType::UndergroundBelt(_, underground_type) => match underground_type {
                UndergroundType::Input => Some(String::from("input")),
                UndergroundType::Output => Some(String::from("output")),
            },
            _ => None,
        }
    }

    pub fn get_maybe_override_stack_size(&self) -> Option<u8> {
        match &self.entity {
            EntityType::Inserter {
                inserter_type: _,
                override_stack_size,
            } => *override_stack_size,
            _ => None,
        }
    }

    pub fn size(&self) -> Position {
        match &self.entity {
            EntityType::Splitter(_) => Position::new(2, 1),
            EntityType::Unknown {
                name: _,
                size,
                misc: _,
            } => *size,
            EntityType::ElectricPole(electric_pole_type) => electric_pole_type.size(),
            _ => Position::new(1, 1),
        }
    }

    pub fn transform(&mut self, transform: Transformation) {
        self.position = self.position.transform(transform);
        self.direction = self.direction.transform(transform);
    }
}

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

static ENTITY_COUNTER: AtomicUsize = AtomicUsize::new(0);

pub struct Blueprint {
    entities: Vec<(EntityKey, Entity)>,
    keys: HashMap<EntityKey, usize>,
    wires: Vec<(EntityKey, u8, EntityKey, u8)>,
}

impl Blueprint {
    pub fn new() -> Self {
        Self {
            entities: Vec::new(),
            keys: HashMap::new(),
            wires: Vec::new(),
        }
    }

    pub fn add_entity(&mut self, entity: Entity) -> EntityKey {
        let id = EntityKey(ENTITY_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed));
        let pos = self.entities.len();

        self.keys.insert(id, pos);
        self.entities.push((id, entity));

        id
    }

    pub fn add_path<'a>(
        &mut self,
        path: impl IntoIterator<Item = &'a PathField>,
        beltspeed: Beltspeed,
    ) {
        for &p in path {
            match p {
                PathField::Belt { pos, dir } => {
                    self.add_entity(Entity::new_belt(beltspeed, pos, dir));
                }
                PathField::Underground { pos, dir, len } => {
                    self.add_entity(Entity::new_underground_belt(
                        beltspeed,
                        UndergroundType::Input,
                        pos,
                        dir,
                    ));
                    self.add_entity(Entity::new_underground_belt(
                        beltspeed,
                        UndergroundType::Output,
                        pos.in_direction(&dir, len as PositionType),
                        dir,
                    ));
                }
            };
        }
    }

    pub fn add_wire(&mut self, a: EntityKey, endpoint_a: u8, b: EntityKey, endpoint_b: u8) {
        self.wires.push((a, endpoint_a, b, endpoint_b));
    }

    pub fn add_blueprint(&mut self, other: Self) {
        self.entities.extend(other.entities);
        self.keys.extend(other.keys);
        self.wires.extend(other.wires);
    }

    pub fn to_blueprint(&self) -> super::Blueprint {
        let entities = self
            .entities
            .iter()
            .enumerate()
            .map(|(i, (_, e))| {
                let size = (e.size() - Position::new(1, 1))
                    .transform(Transformation::new(e.direction, Position::new(0, 0)));
                BlueprintEntity::builder(
                    e.get_name(),
                    i as u32 + 1,
                    BlueprintPosition::new(
                        e.position.x as f64 + 0.5 * size.x as f64 + 0.5,
                        e.position.y as f64 + 0.5 * size.y as f64 + 0.5,
                    ),
                )
                .direction(match e.direction {
                    Direction::Up => 0,
                    Direction::Right => 4,
                    Direction::Down => 8,
                    Direction::Left => 12,
                })
                .maybe_underground_type(e.get_maybe_underground_type_string())
                .maybe_override_stack_size(e.get_maybe_override_stack_size())
                .build()
            })
            .collect();

        let wires = self
            .wires
            .iter()
            .map(|&(a, endpoint_a, b, endpoint_b)| {
                [
                    self.keys[&a] as u32 + 1,
                    endpoint_a as u32,
                    self.keys[&b] as u32 + 1,
                    endpoint_b as u32,
                ]
            })
            .collect();

        super::Blueprint::builder()
            .label(String::from("test"))
            .entities(entities)
            .wires(wires)
            .build()
    }

    pub fn transform(&mut self, transform: Transformation) {
        for (_, e) in &mut self.entities {
            e.transform(transform);
        }
    }
}

impl Default for Blueprint {
    fn default() -> Self {
        Self::new()
    }
}