use egui_snarl::{InPinId, NodeId, OutPinId, Snarl};
use std::array::from_fn;
use std::fmt;
use topological_sort::TopologicalSort;

use crate::node::{AnyNode, CompileError, ConcreteNode, MAX_INPUTS};
use crate::types::{Type, TypeSignature};

pub trait WithSignatures {
    // internal
    fn get_input_types(&self, node: NodeId) -> [Option<Type>; MAX_INPUTS];

    // for rendering
    fn get_visible_inputs(&self, node: NodeId) -> Vec<Option<Type>>;

    // for compiling
    fn get_node_signature(&self, node: NodeId) -> Result<TypeSignature, CompileError>;
    fn get_node_signature_with_overrides(
        &self,
        node: NodeId,
        overrides: &[Option<Type>; MAX_INPUTS],
    ) -> Result<TypeSignature, CompileError>;
}

impl WithSignatures for Snarl<AnyNode> {
    fn get_input_types(&self, node: NodeId) -> [Option<Type>; MAX_INPUTS] {
        from_fn(
            move |input| match &*self.in_pin(InPinId { node, input }).remotes {
                [] => None,
                [out_pin] => self
                    .get_node_signature(out_pin.node)
                    .map(|s| s.outputs[out_pin.output])
                    .ok(),
                _ => unreachable!("cannot connect to multiple inputs"),
            },
        )
    }

    fn get_visible_inputs(&self, node: NodeId) -> Vec<Option<Type>> {
        self[node].visible_inputs(&self.get_input_types(node))
    }

    fn get_node_signature(&self, node: NodeId) -> Result<TypeSignature, CompileError> {
        self[node].signature(&self.get_input_types(node))
    }

    fn get_node_signature_with_overrides(
        &self,
        node: NodeId,
        overrides: &[Option<Type>; MAX_INPUTS],
    ) -> Result<TypeSignature, CompileError> {
        let mut inputs = self.get_input_types(node);
        for (i, over) in overrides.iter().enumerate() {
            inputs[i] = over.or(inputs[i]);
        }
        self[node].signature(&inputs)
    }
}

#[derive(fmt::Debug, Copy, Clone)]
pub struct GraphError(pub CompileError, pub Option<NodeId>);
impl From<fmt::Error> for GraphError {
    fn from(v: fmt::Error) -> Self {
        Self(v.into(), None)
    }
}

pub trait Compilable {
    fn compile_node(&self, f: &mut dyn fmt::Write, node: NodeId) -> Result<(), CompileError>;
    fn compile(&self, f: &mut dyn fmt::Write) -> Result<(), GraphError>;
}

fn compile_output(pin: OutPinId) -> String {
    let node_id = pin.node.0;
    let out_id = pin.output;
    format!("n{node_id}_o{out_id}")
}

impl Compilable for Snarl<AnyNode> {
    fn compile(&self, f: &mut dyn fmt::Write) -> Result<(), GraphError> {
        let mut order = TopologicalSort::<NodeId>::new();
        for (out, inp) in self.wires() {
            order.add_dependency(out.node, inp.node);
        }

        write!(
            f,
            "
#version 450
in vec2 in_uv;
out vec4 out_color;

layout(std140, binding = 0) uniform Uniforms {{
    vec2 in_resolution;
    float in_time;
    float _pad3;
}} uniforms;

void main() {{
"
        )?;
        while let Some(id) = order.pop() {
            self.compile_node(f, id)
                .map_err(|inner| GraphError(inner, Some(id)))?;
        }
        writeln!(f, "}}")?;

        Ok(())
    }

    fn compile_node(&self, f: &mut dyn fmt::Write, node: NodeId) -> Result<(), CompileError> {
        let signature = self.get_node_signature(node)?;

        // log::info!("{node} {self[node]}: {signature:?}");

        let inputs: Vec<Option<String>> = (0..signature.inputs.len())
            .map(
                |input| match &*self.in_pin(InPinId { node, input }).remotes {
                    [] => None,
                    [pin] => Some(compile_output(*pin)),
                    _ => unreachable!("cannot connect to multiple inputs"),
                },
            )
            .collect();

        let outputs: Vec<String> = (0..signature.outputs.len())
            .map(|output| compile_output(OutPinId { node, output }))
            .collect();

        self[node].compile(f, signature, inputs, outputs)
    }
}