use prost::Message;
use crate::tfpb::tensorflow::{GraphDef, NodeDef, SavedModel};
use std::{fs, path};
use tract_core::internal::*;

#[derive(Default)]
pub struct ParsingContext {
    pub node_output_arities: HashMap<String, usize>
}

#[derive(Clone, Default)]
pub struct TfOpRegister(
    pub HashMap<String, fn(&ParsingContext, node: &NodeDef) -> TractResult<Box<dyn InferenceOp>>>,
);

impl TfOpRegister {
    pub fn insert(
        &mut self,
        s: &'static str,
        builder: fn(&ParsingContext, node: &NodeDef) -> TractResult<Box<dyn InferenceOp>>,
    ) {
        self.0.insert(s.into(), builder);
    }
}

pub struct Tensorflow {
    pub op_register: TfOpRegister,
}

impl Tensorflow {
    // From the node_def.proto documentation:
    // Each input is "node:src_output" with "node" being a string name and
    // "src_output" indicating which output tensor to use from "node". If
    // "src_output" is 0 the ":0" suffix can be omitted. Regular inputs may
    // optionally be followed by control inputs that have the format "^node".
    fn parse_input(i: &str) -> TractResult<(&str, usize)> {
        let pair = if i.starts_with("^") {
            (&i[1..], 0)
        } else {
            let splits: Vec<_> = i.splitn(2, ':').collect();
            (splits[0], if splits.len() > 1 { splits[1].parse::<usize>()? } else { 0 })
        };
        Ok(pair)
    }

    pub fn determinize(model: &mut GraphDef) -> TractResult<()> {
        for pbnode in &mut model.node {
            if pbnode.op == "RandomUniform" {
                if pbnode.get_attr_int::<i64>("seed")? == 0
                    && pbnode.get_attr_int::<i64>("seed2")? == 0
                {
                    pbnode.attr.insert("seed".to_string(), 1.into());
                    pbnode.attr.insert("seed2".to_string(), 1.into());
                }
            }
        }
        Ok(())
    }

    pub fn read_frozen_model(&self, r: &mut dyn std::io::Read) -> TractResult<GraphDef> {
        let mut v = vec!();
        r.read_to_end(&mut v)?;
        let b = bytes::Bytes::from(v);
        Ok(GraphDef::decode(b).map_err(|e| format!("{:?}", e))?)
    }

    pub fn open_saved_model(&self, r: &mut dyn std::io::Read) -> TractResult<SavedModel> {
        let mut v = vec!();
        r.read_to_end(&mut v)?;
        let b = bytes::Bytes::from(v);
        Ok(SavedModel::decode(b).map_err(|e| format!("{:?}", e))?)
    }

    /// Convenience method: will read the first model in the saved model
    /// container. Use open_avec_model for more control.
    pub fn read_saved_model(&self, r: &mut dyn std::io::Read) -> TractResult<GraphDef> {
        let mut saved = self.open_saved_model(r)?;
        Ok(saved.meta_graphs.remove(0).graph_def.unwrap())
    }
}

impl Framework<GraphDef> for Tensorflow {
    /// This method will try to read as frozen model, then as a saved model.
    fn proto_model_for_path(&self, r: impl AsRef<path::Path>) -> TractResult<GraphDef> {
        self.read_frozen_model(&mut fs::File::open(r.as_ref())?)
            .or_else(|_| self.read_saved_model(&mut fs::File::open(r.as_ref())?))
    }

    /// This method expects a frozen model, use open_saved_model for TF2 saved
    /// model format.
    fn proto_model_for_read(&self, r: &mut dyn std::io::Read) -> TractResult<GraphDef> {
        self.read_frozen_model(r)
    }

    fn model_for_proto_model(&self, graph: &GraphDef) -> TractResult<InferenceModel> {
        use crate::ops::control_flow as cf;

        let mut model = InferenceModel::default();
        let mut inputs = tvec!();
        let mut context = ParsingContext::default();

        // compute min output arity for all nodes
        for pbnode in &graph.node {
            for i in &pbnode.input {
                let (node, slot) = Self::parse_input(i)?;
                let arity = context.node_output_arities.entry(node.to_string()).or_insert(1);
                *arity = (*arity).max(slot + 1);
            }
        }

        for pbnode in &graph.node {
            let name = &pbnode.name;

            if pbnode.op == "NextIteration" {
                let source_op = cf::NextIteration::new(name.clone(), cf::NextIterationRole::Source);
                let sink_op = cf::NextIteration::new(name.clone(), cf::NextIterationRole::Sink);
                let _source =
                    model.add_node(name.clone(), source_op, tvec!(InferenceFact::default()))?;
                let _sink = model.add_node(format!("{}-Sink", name), sink_op, tvec!())?;
                continue;
            }

            let op = match self.op_register.0.get(&pbnode.op) {
                Some(builder) => (builder)(&context, pbnode)?,
                None => tract_core::ops::unimpl::UnimplementedOp::new(
                    context.node_output_arities.get(name).cloned().unwrap_or(1),
                    &pbnode.op,
                    format!("{:?}", pbnode),
                )
                .into(),
            };

            let facts = tvec!(InferenceFact::default(); op.nboutputs()?);

            let node_id = model.add_node(name.clone(), op, facts)?;
            if pbnode.op == "Placeholder" {
                let dt = pbnode.get_attr_datum_type("dtype")?;
                let mut fact = InferenceFact::dt(dt);
                if let Some(shape) = pbnode.get_attr_opt_shape("shape")? {
                    let shape_fact = ShapeFact::closed(
                        shape
                            .iter()
                            .map(|d| {
                                if *d == -1 {
                                    GenericFact::Any
                                } else {
                                    GenericFact::Only(d.to_dim())
                                }
                            })
                            .collect(),
                    );
                    fact = fact.with_shape(shape_fact);
                }
                inputs.push(OutletId::new(node_id, 0));
                model.set_outlet_fact(OutletId::new(node_id, 0), fact)?;
            }
        }

        for pbnode in &graph.node {
            let node_id = if pbnode.op == "NextIteration" {
                model.node_by_name(&*format!("{}-Sink", &pbnode.name))?.id
            } else {
                model.node_by_name(&pbnode.name)?.id
            };
            for (ix, i) in pbnode.input.iter().enumerate() {
                let input = Self::parse_input(i)?;
                let prec = model.node_by_name(input.0)?.id;
                if i.starts_with("^") {
                    panic!();
                } else {
                    let outlet = OutletId::new(prec, input.1);
                    let inlet = InletId::new(node_id, ix);
                    model.add_edge(outlet, inlet)?;
                    model.set_outlet_label(outlet, i.to_string());
                }
            }
        }

        // variable -> assign rewire
        //  * Assign consumes this by_ref tensor on #0 and somehow performs
        //      updates on it (it has a second input on #1 for the value to
        //      assign)
        //
        // in tract:
        //  * VariableV2 outputs a regular tensor stored in the session state
        //  * Assign has the same inputs, but do not uses the #0, udating the
        //      state session instead
        for id in 0..model.nodes().len() {
            use crate::ops::vars::*;
            if model.node(id).op_is::<Assign>() {
                let prec = model.node(id).inputs[0];
                let var_id = model.node(prec.node).op_as::<VariableV2>().map(|v| v.id.clone());
                if let (Some(var_id), Some(assign)) =
                    (var_id, model.node_mut(id).op_as_mut::<Assign>())
                {
                    assign.var_id = Some(var_id);
                } else {
                    bail!("Model contains unlinked Assign/Variable2");
                }
            }
        }
        model.set_input_outlets(&*inputs)?;
        model.auto_outputs()?;
        Ok(model)
    }
}