use crate::migration::fields::FieldDescriptor;
use crate::migration::registry::{MigrationDependency, MigrationStepEntry};
use crate::migration::set::MigrationSet;
use crate::{MigrationContext, Migrator, StateScope};
use std::collections::{BTreeSet, HashMap, HashSet};

#[derive(Default)]
pub struct MigrationBuilder {
    prefixes: HashMap<String, PrefixPlan>,
}

#[derive(Default)]
pub(crate) struct PrefixPlan {
    migrator: Migrator,
    dependencies: BTreeSet<String>,
    pub(crate) fields: &'static [FieldDescriptor],
    pub(crate) schema_hash: u64,
}

pub struct PrefixMigrationBuilder<'a> {
    builder: &'a mut MigrationBuilder,
    prefix: String,
}

impl MigrationBuilder {
    pub fn collect_codegen(&mut self) -> &mut Self {
        let mut groups: HashMap<&'static str, Vec<&'static MigrationStepEntry>> = HashMap::new();

        for entry in inventory::iter::<MigrationStepEntry> {
            groups.entry(entry.prefix).or_default().push(entry);
        }

        for (prefix, steps) in groups {
            let mut latest_hash = 0;
            let mut max_v = 0;
            let mut latest_fields: &'static [FieldDescriptor] = &[];
            let mut merged_deps = HashSet::new();

            for step in &steps {
                if step.target_version >= max_v {
                    max_v = step.target_version;
                    latest_hash = step.schema_hash;
                    latest_fields = step.fields;
                }

                for dep in step.dependencies {
                    merged_deps.insert(*dep);
                }

                if step.target_version > 0 {
                    self.for_prefix(prefix)
                        .step(step.target_version, step.description, step.run);
                }
            }

            let plan = self.prefix_plan(prefix);
            plan.schema_hash = latest_hash;
            plan.fields = latest_fields;
            for dep in merged_deps {
                plan.dependencies.insert(dep.to_string());
            }
        }
        self
    }

    pub fn for_node<T: StateScope>(&mut self) -> PrefixMigrationBuilder<'_> {
        self.for_prefix(T::PREFIX)
    }

    pub fn for_prefix(&mut self, prefix: impl Into<String>) -> PrefixMigrationBuilder<'_> {
        PrefixMigrationBuilder {
            builder: self,
            prefix: prefix.into(),
        }
    }

    pub(crate) fn prefix_plan(&mut self, prefix: &str) -> &mut PrefixPlan {
        self.prefixes.entry(prefix.to_string()).or_default()
    }

    pub(crate) fn into_set(self) -> MigrationSet {
        let mut set = MigrationSet::default();
        let mut prefixes = self.prefixes.into_iter().collect::<Vec<_>>();

        prefixes.sort_by(|(a, _), (b, _)| a.cmp(b));

        for (prefix, plan) in prefixes {
            let deps: Vec<&str> = plan.dependencies.iter().map(|s| s.as_str()).collect();
            set = set.add(prefix, plan.migrator, plan.schema_hash, plan.fields, &deps);
        }
        set
    }
}

impl PrefixMigrationBuilder<'_> {
    pub fn depends_on<D: MigrationDependency>(&mut self) -> &mut Self {
        let plan = self.builder.prefix_plan(&self.prefix);
        D::register(&mut plan.dependencies);
        self
    }

    pub fn depends_on_raw(&mut self, dependency: impl Into<String>) -> &mut Self {
        self.builder
            .prefix_plan(&self.prefix)
            .dependencies
            .insert(dependency.into());
        self
    }

    pub fn step<F>(&mut self, target_version: u32, description: &str, run: F) -> &mut Self
    where
        F: Fn(&mut MigrationContext) -> crate::Result<()> + Send + Sync + 'static,
    {
        let plan = self.builder.prefix_plan(&self.prefix);
        let migrator = std::mem::take(&mut plan.migrator);
        plan.migrator = migrator.step(target_version, description, run);
        self
    }
}
