// SPDX-FileCopyrightText: 2024 Topola contributors
//
// SPDX-License-Identifier: MIT

use std::collections::BTreeMap;

use enum_dispatch::enum_dispatch;
use geo::Point;
use petgraph::{
    data::Element,
    graph::NodeIndex,
    prelude::StableUnGraph,
    unionfind::UnionFind,
    visit::{EdgeRef, IntoEdgeReferences, NodeIndexable},
};
use spade::{HasPosition, InsertionError, Point2};

use crate::{
    drawing::{
        band::BandTermsegIndex,
        dot::FixedDotIndex,
        graph::{GetMaybeNet, MakePrimitiveRef, PrimitiveIndex},
        primitive::MakePrimitiveShape,
        rules::AccessRules,
    },
    geometry::shape::AccessShape,
    graph::{GenericIndex, GetPetgraphIndex, MakeRef},
    layout::{
        poly::{MakePolygon, PolyWeight},
        Layout,
    },
    triangulation::{GetTrianvertexNodeIndex, Triangulation},
};

use super::ratline::{RatlineIndex, RatlineWeight};

#[enum_dispatch(GetPetgraphIndex)]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum RatvertexIndex {
    FixedDot(FixedDotIndex),
    Poly(GenericIndex<PolyWeight>),
}

impl From<RatvertexIndex> for crate::layout::NodeIndex {
    fn from(vertex: RatvertexIndex) -> crate::layout::NodeIndex {
        match vertex {
            RatvertexIndex::FixedDot(dot) => crate::layout::NodeIndex::Primitive(dot.into()),
            RatvertexIndex::Poly(poly) => crate::layout::NodeIndex::Compound(poly.into()),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct RatvertexWeight {
    vertex: RatvertexIndex,
    pub pos: Point,
}

impl GetTrianvertexNodeIndex<RatvertexIndex> for RatvertexWeight {
    fn node_index(&self) -> RatvertexIndex {
        self.vertex
    }
}

impl HasPosition for RatvertexWeight {
    type Scalar = f64;
    fn position(&self) -> Point2<Self::Scalar> {
        Point2::new(self.pos.x(), self.pos.y())
    }
}

pub struct Ratsnest {
    graph: StableUnGraph<RatvertexWeight, RatlineWeight, usize>,
}

impl Ratsnest {
    pub fn new(layout: &Layout<impl AccessRules>) -> Result<Self, InsertionError> {
        let mut unionfind = UnionFind::new(layout.drawing().geometry().graph().node_bound());

        for edge in layout.drawing().geometry().graph().edge_references() {
            unionfind.union(edge.source(), edge.target());
        }

        let mut this = Self {
            graph: StableUnGraph::default(),
        };

        let mut triangulations = BTreeMap::new();
        let node_bound = layout.drawing().geometry().graph().node_bound();

        for layer in 0..layout.drawing().layer_count() {
            let mut handle_ratvertex_weight =
                |maybe_net: Option<usize>, vertex: RatvertexIndex, pos: Point| {
                    if let Some(net) = maybe_net {
                        triangulations
                            .entry((layer, net))
                            .or_insert_with(|| Triangulation::new(node_bound))
                            .add_vertex(RatvertexWeight { vertex, pos })?;
                    }
                    Ok(())
                };

            for node in layout.drawing().layer_primitive_nodes(layer) {
                if let PrimitiveIndex::FixedDot(dot) = node {
                    // Dots that are parts of polys are ignored because ratlines
                    // should only go to their centerpoints.
                    if layout.drawing().compounds(dot).next().is_none() {
                        handle_ratvertex_weight(
                            layout.drawing().primitive(dot).maybe_net(),
                            RatvertexIndex::FixedDot(dot),
                            node.primitive_ref(layout.drawing()).shape().center(),
                        )?;
                    }
                }
            }

            for poly in layout.layer_poly_nodes(layer) {
                handle_ratvertex_weight(
                    layout.drawing().compound_weight(poly.into()).maybe_net(),
                    RatvertexIndex::Poly(poly),
                    poly.ref_(layout).shape().center(),
                )?;
            }
        }

        for ((_layer, _net), triangulation) in triangulations {
            let mut map = Vec::new();

            for element in petgraph::algo::min_spanning_tree(&triangulation) {
                match element {
                    Element::Node { weight } => {
                        map.push(this.graph.add_node(weight));
                    }
                    Element::Edge {
                        source,
                        target,
                        weight,
                    } => {
                        this.graph.add_edge(map[source], map[target], weight.weight);
                    }
                }
            }
        }

        this.graph.retain_edges(|g, i| {
            if let Some((source, target)) = g.edge_endpoints(i) {
                let source_index = g.node_weight(source).unwrap().node_index().petgraph_index();
                let target_index = g.node_weight(target).unwrap().node_index().petgraph_index();
                !unionfind.equiv(source_index, target_index)
            } else {
                true
            }
        });

        Ok(this)
    }

    pub fn assign_band_termseg_to_ratline(
        &mut self,
        ratline: RatlineIndex,
        termseg: BandTermsegIndex,
    ) {
        self.graph.edge_weight_mut(ratline).unwrap().band_termseg = Some(termseg);
    }

    pub fn graph(&self) -> &StableUnGraph<RatvertexWeight, RatlineWeight, usize> {
        &self.graph
    }
}