use std::{collections::HashMap, fmt, iter};

use crate::{
    data::{self, Path},
    exec::{self, ExecNode, PreparedRoot, ResolveError, Tree},
    syn, Expr, Ident, Shaped, Value, R,
};

pub type TypeExpr = crate::TypeExpr<NodeId, Expr<R<Ident>>>;
pub type Field = crate::Field<TypeExpr>;
pub type NodeApp = crate::NodeApp<NodeId, Expr<R<Ident>>>;

#[derive(Debug)]
pub struct Node {
    properties: HashMap<String, lexpr::Value>,
    field_ids: Vec<FieldId>,
    implications: Vec<NodeApp>,
    branches: Vec<Branch>,
}

impl Node {
    fn implications(&self) -> &[NodeApp] {
        &self.implications
    }

    fn field_ids(&self) -> &[FieldId] {
        &self.field_ids
    }

    fn find_branch(&self, name: &Ident) -> Option<usize> {
        self.branches
            .iter()
            .enumerate()
            .find_map(|(i, b)| if &b.name == name { Some(i) } else { None })
    }
}

#[derive(Debug)]
struct Branch {
    name: Ident,
    target: NodeId,
    condition: data::Expr,
}

impl Branch {
    fn name(&self) -> &Ident {
        &self.name
    }
}

#[derive(Eq, PartialEq, Debug, Copy, Clone)]
struct FieldId(usize);

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

#[derive(Debug)]
pub struct Error(InnerError);

impl Error {
    pub fn display<'a>(&'a self, forest: &'a Forest) -> ErrorDisplay<'a> {
        ErrorDisplay {
            error: self,
            forest,
        }
    }

    fn no_such_root(name: &Ident) -> Error {
        Error(InnerError::NoSuchRoot(name.clone()))
    }

    fn no_such_branch(
        root_name: &Ident,
        root_id: NodeId,
        branch_ids: Vec<usize>,
        name: &Ident,
    ) -> Error {
        Error(InnerError::NoSuchBranch {
            root_name: root_name.clone(),
            root_id,
            branch_ids,
            name: name.clone(),
        })
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use InnerError::*;
        match &self.0 {
            NoSuchRoot(name) => write!(f, "no such root: '{}'", name),
            NoSuchBranch {
                root_name, name, ..
            } => write!(f, "unknown branch below {}: {}", root_name, name),
        }
    }
}

#[derive(Debug)]
enum InnerError {
    NoSuchRoot(Ident),
    NoSuchBranch {
        root_name: Ident,
        root_id: NodeId,
        branch_ids: Vec<usize>,
        name: Ident,
    },
}

pub struct ErrorDisplay<'a> {
    #[allow(dead_code)] // See FIXME below
    forest: &'a Forest,
    error: &'a Error,
}

impl<'a> fmt::Display for ErrorDisplay<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use InnerError::*;
        match &self.error.0 {
            NoSuchRoot(name) => write!(f, "no such root: '{}'", name),
            NoSuchBranch {
                root_name, name, ..
            } => {
                // FIXME: display full path
                write!(f, "unknown branch below {}: {}", root_name, name)
            }
        }
    }
}

#[derive(Debug)]
pub struct PrepareError(InnerPrepareError);

#[derive(Debug)]
enum InnerPrepareError {
    NoSuchRoot(Ident),
    Resolution(exec::ResolveError),
}

impl PrepareError {
    fn no_such_root(name: Ident) -> Self {
        PrepareError(InnerPrepareError::NoSuchRoot(name))
    }
    fn resolution(e: exec::ResolveError) -> Self {
        PrepareError(InnerPrepareError::Resolution(e))
    }
}

impl fmt::Display for PrepareError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use InnerPrepareError::*;
        match &self.0 {
            NoSuchRoot(name) => write!(f, "root '{}' does not exist", name),
            Resolution(e) => write!(f, "error in name resolution: {}", e),
        }
    }
}

impl std::error::Error for PrepareError {}

#[derive(Debug, Default)]
pub struct Forest {
    nodes: Vec<Node>,
    fields: Vec<Field>,
    roots: HashMap<Ident, NodeId>,
    root_names: Vec<Ident>, // in insertion order
}

#[derive(Debug, Clone)]
pub struct Cursor<'a> {
    forest: &'a Forest,
    tree: &'a exec::Tree,
    node: &'a ExecNode,
    node_fields: FieldIter<'a>,
    branch_ids: &'a [usize],
    decoded: &'a [exec::DecodedItem],
    skip_hidden: bool,
}

impl<'a> Cursor<'a> {
    fn is_empty(&self) -> bool {
        if self.skip_hidden {
            self.clone().next().is_none()
        } else {
            self.decoded.is_empty()
        }
    }
}

impl<'a> Iterator for Cursor<'a> {
    type Item = (&'a Field, &'a Value);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(field) = self.node_fields.next() {
                let (item, rest) = self.decoded.split_first().expect("invalid cursor");
                self.decoded = rest;
                if self.skip_hidden && field.hidden {
                    continue;
                }
                return Some((field, item.value()));
            } else if let Some((branch_id, branch_ids)) = self.branch_ids.split_first() {
                self.node = &self.tree[self.node.branches()[*branch_id].target];
                self.node_fields = self.forest.fields(self.node.node_id());
                self.branch_ids = branch_ids;
            } else {
                return None;
            }
        }
    }
}

#[derive(Debug)]
pub struct Record<'a> {
    forest: &'a Forest,
    tree: &'a exec::Tree,
    node: &'a ExecNode,
    branch_ids: &'a [usize],
    decoded: &'a [exec::DecodedItem],
}

impl<'a> Record<'a> {
    pub fn new(shaped: &'a Shaped, forest: &'a Forest, tree: &'a Tree) -> Self {
        Record {
            forest,
            tree,
            node: &tree[shaped.root_id()],
            branch_ids: shaped.branch_ids(),
            decoded: shaped.decoded(),
        }
    }

    fn cursor(&'a self, skip_hidden: bool) -> Cursor<'a> {
        Cursor {
            forest: self.forest,
            tree: self.tree,
            node: self.node,
            node_fields: self.forest.fields(self.node.node_id()),
            branch_ids: &self.branch_ids,
            decoded: self.decoded,
            skip_hidden,
        }
    }
}

impl<'a> fmt::Display for Record<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut branches = self.forest.follow(self.node.node_id(), &self.branch_ids);
        let anonymous = branches.is_empty();
        while let Some(branch) = branches.next() {
            f.write_str(branch.name().as_str())?;
            if !branches.is_empty() {
                f.write_str(".")?;
            }
        }
        let mut cursor = self.cursor(true);
        if cursor.is_empty() {
            return Ok(());
        }
        if anonymous {
            f.write_str("{ ")?;
        } else {
            f.write_str(" { ")?;
        }
        while let Some((field, value)) = cursor.next() {
            write!(
                f,
                "{}: {}",
                field.name,
                value.display(self.forest, self.tree)
            )?;
            if !cursor.is_empty() {
                f.write_str(", ")?;
            }
        }
        f.write_str(" }")
    }
}

impl<'a> From<Record<'a>> for lexpr::Value {
    fn from(record: Record<'a>) -> Self {
        lexpr::Value::from(&record)
    }
}

impl<'a> From<&Record<'a>> for lexpr::Value {
    fn from(record: &Record<'a>) -> Self {
        use lexpr::Value;
        let branch_names: Vec<_> = record
            .forest
            .follow(record.node.node_id(), &record.branch_ids)
            .map(|branch| branch.name().as_str())
            .collect();
        let fields = record.cursor(true).map(|(field, value)| {
            Value::cons(
                Value::symbol(field.name.as_str()),
                value.to_sexp_named(record.forest, record.tree),
            )
        });
        if branch_names.is_empty() {
            Value::list(fields)
        } else {
            let name = branch_names.join(".");
            Value::list(iter::once(Value::symbol(name)).chain(fields))
        }
    }
}

/// Iterator for the fields of a node, including implied ones
#[derive(Debug, Clone)]
struct FieldIter<'a> {
    forest: &'a Forest,
    implied_iter: Option<Box<FieldIter<'a>>>,
    implications: std::slice::Iter<'a, NodeApp>,
    field_ids: std::slice::Iter<'a, FieldId>,
}

impl<'a> FieldIter<'a> {
    fn new(forest: &'a Forest, node_id: NodeId) -> Self {
        FieldIter {
            forest,
            implied_iter: None,
            implications: forest[node_id].implications().iter(),
            field_ids: forest[node_id].field_ids().iter(),
        }
    }
}

impl<'a> std::iter::Iterator for FieldIter<'a> {
    type Item = &'a Field;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(implied) = &mut self.implied_iter {
                if let Some(implied_field) = implied.next() {
                    return Some(implied_field);
                }
                self.implied_iter = None;
            }
            if let Some(implication) = self.implications.next() {
                self.implied_iter = Some(Box::new(self.forest.fields(implication.id)));
            } else {
                break;
            }
        }
        self.field_ids
            .next()
            .map(|field_id| &self.forest.fields[field_id.0])
    }
}

impl Forest {
    pub fn new() -> Self {
        Forest::default()
    }

    pub fn root_names(&self) -> std::slice::Iter<'_, Ident> {
        self.root_names.iter()
    }

    pub fn add_root(&mut self, root: data::Root) -> Result<(), Error> {
        let node = Node {
            properties: Default::default(),
            implications: Default::default(),
            field_ids: self.push_fields(root.fields)?,
            branches: Default::default(),
        };
        let node_id = self.push_node(node);
        self.root_names.push(root.name.clone());
        self.roots.insert(root.name, node_id);
        Ok(())
    }

    pub fn add_branch(&mut self, branch: data::Branch) -> Result<(), Error> {
        let root_id = self.named_root_id(&branch.root)?;
        let (parent_id, _trace) = self.trace(root_id, &branch.path, &branch.root)?;
        let implications = branch
            .implies
            .into_iter()
            .map(|app| self.push_node_app(app))
            .collect::<Result<_, _>>()?;
        // FIXME: check for duplicate branch
        let node = Node {
            properties: branch.properties,
            implications,
            field_ids: self.push_fields(branch.fields)?,
            branches: Default::default(),
        };
        let node_id = self.push_node(node);
        let parent = &mut self.nodes[parent_id.0];
        parent.branches.push(Branch {
            name: branch.name,
            target: node_id,
            condition: branch.condition,
        });
        Ok(())
    }

    // In case of error, the enum may be added partially to the forest.
    pub fn add_enum(&mut self, en: data::Enum) -> Result<(), Error> {
        let mut en = en;
        let node = Node {
            properties: Default::default(),
            implications: Vec::new(), // TODO: add support
            field_ids: self.push_fields(en.discriminants)?,
            branches: Vec::new(),
        };
        let node_id = self.push_node(node);
        self.root_names.push(en.name.clone());
        self.roots.insert(en.name, node_id);
        let mut branches = Vec::new();
        for variant in en.variants.drain(0..) {
            let target = Node {
                properties: variant.properties,
                implications: Vec::new(), // TODO: add support
                field_ids: self.push_fields(variant.fields)?,
                branches: Vec::new(),
            };
            let target_id = self.push_node(target);
            branches.push(Branch {
                name: variant.name,
                target: target_id,
                condition: variant.condition,
            });
        }
        self.nodes[node_id.0].branches = branches;
        Ok(())
    }

    pub fn nodes(&self) -> std::slice::Iter<'_, Node> {
        self.nodes.iter()
    }

    fn follow<'a>(&'a self, node_id: NodeId, branch_ids: &'a [usize]) -> BranchIter<'a> {
        let node = &self[node_id];
        BranchIter {
            forest: self,
            node,
            branch_ids,
        }
    }

    fn trace(
        &self,
        root_id: NodeId,
        path: &Path,
        root_name: &Ident,
    ) -> Result<(NodeId, Vec<usize>), Error> {
        let mut node_id = root_id;
        let mut trace = Vec::with_capacity(path.len());
        for ident in path {
            let node = &self[node_id];
            let branch_idx = match node.find_branch(ident) {
                Some(idx) => idx,
                None => return Err(Error::no_such_branch(root_name, root_id, trace, ident)),
            };
            trace.push(branch_idx);
            node_id = node.branches[branch_idx].target;
        }
        Ok((node_id, trace))
    }

    fn fields(&self, node_id: NodeId) -> FieldIter<'_> {
        FieldIter::new(self, node_id)
    }

    // Will panic if trace contains invalid references. This should only be
    // possible if using a trace from a different forest, as all traces that can
    // be produced should be (and stay) valid for that forest.
    pub fn resolve<'a>(&'a self, prepared: &'a PreparedRoot, shaped: &'a Shaped) -> Record<'a> {
        let tree = prepared.tree();
        Record {
            forest: self,
            tree,
            node: &tree[shaped.root_id()],
            branch_ids: shaped.branch_ids(),
            decoded: shaped.decoded(),
        }
    }

    pub fn prepare(&self, root: &Ident) -> Result<PreparedRoot, PrepareError> {
        let root_id = *self
            .roots
            .get(root)
            .ok_or_else(|| PrepareError::no_such_root(root.clone()))?;
        let mut prepare = exec::PrepareContext::new();
        let exec_root =
            self.prepare_node(root_id, root_id, &mut prepare, &mut exec::FieldEnv::new())?;
        Ok(prepare.finish(exec_root))
    }

    // TODO: Move `shapes` and `fields` into a new `Context` struct
    fn prepare_node(
        &self,
        root_id: NodeId,
        node_id: NodeId,
        ctx: &mut exec::PrepareContext,
        field_env: &mut exec::FieldEnv,
    ) -> Result<exec::ExecNodeId, PrepareError> {
        let node = &self.nodes[node_id.0];
        let mut implications = Vec::with_capacity(node.implications.len());
        let mut app_state = (ctx, Vec::new());
        for app in &node.implications {
            let lifted_app = app.lift(
                &mut app_state,
                |node_id, (ctx, sub_envs)| {
                    let node_id = *node_id;
                    let exec_id = ctx.include(node_id, |ctx| {
                        // Implications have an isolated scope
                        let mut sub_env = exec::FieldEnv::new();
                        let prepared = self.prepare_node(node_id, node_id, ctx, &mut sub_env)?;
                        sub_envs.push(sub_env);
                        Ok(prepared)
                    })?;
                    Ok(exec_id)
                },
                |expr, _ctx| resolve_expr(expr, field_env).map_err(PrepareError::resolution),
            )?;
            implications.push(lifted_app);
        }
        let (ctx, sub_envs) = app_state;
        for sub_env in sub_envs {
            field_env.extend(sub_env);
        }
        // First, resolve all field types, which have access to field names, but not to field types
        field_env.push_names(
            node.field_ids
                .iter()
                .map(|field_id| &self.fields[field_id.0].name),
        );
        let fields: Vec<_> = node
            .field_ids
            .iter()
            .map(|field_id| {
                let field = &self.fields[field_id.0];
                let ty = field.ty.lift(
                    ctx,
                    |node_id, ctx| {
                        let node_id = *node_id;
                        let exec_id = ctx.include(node_id, |ctx| {
                            let mut fields = exec::FieldEnv::new();
                            self.prepare_node(node_id, node_id, ctx, &mut fields)
                        })?;
                        Ok(exec_id)
                    },
                    |expr, _ctx| resolve_expr(expr, field_env).map_err(PrepareError::resolution),
                )?;
                Ok(exec::Field::new(ty, field.constant.clone()))
            })
            .collect::<Result<_, _>>()?;
        let prev_len = field_env.resolve_types(fields.iter().map(|f| f.ty().clone()));

        let is_root = root_id == node_id;
        let exec_id = ctx.push_node(
            ExecNode::new(node_id, node.properties.clone(), implications, fields),
            is_root,
        );
        // All branch conditions and target nodes have access to the node's
        // parsed fields.
        let branches: Vec<_> = node
            .branches
            .iter()
            .map(|branch| {
                let cond =
                    resolve_expr(&branch.condition, field_env).map_err(PrepareError::resolution)?;
                let target = self.prepare_node(root_id, branch.target, ctx, field_env)?;
                Ok(exec::Branch::new(cond, target))
            })
            .collect::<Result<_, _>>()?;
        ctx.set_branches(exec_id, branches);
        field_env.truncate(prev_len);
        Ok(exec_id)
    }

    fn push_node(&mut self, node: Node) -> NodeId {
        let id = self.nodes.len();
        self.nodes.push(node);
        NodeId(id)
    }

    fn named_root_id(&self, name: &Ident) -> Result<NodeId, Error> {
        match self.roots.get(name) {
            Some(node_id) => Ok(*node_id),
            None => Err(Error::no_such_root(&name)),
        }
    }

    fn push_node_app(&mut self, app: syn::NodeApp) -> Result<NodeApp, Error> {
        let node_id = self.named_root_id(&app.id)?;
        Ok(NodeApp {
            id: node_id,
            args: app.args,
        })
    }

    fn push_type(&mut self, ty: syn::TypeExpr) -> Result<TypeExpr, Error> {
        use crate::NodeApp;
        use crate::TypeExpr::*;
        let ty = match ty {
            I(width, endianness) => I(width, endianness),
            U(width, endianness) => U(width, endianness),
            ArrayN(n, et) => {
                let et = self.push_type(*et)?;
                ArrayN(n, Box::new(et))
            }
            ArrayV(et) => {
                let et = self.push_type(*et)?;
                ArrayV(Box::new(et))
            }
            App(app) => {
                let node_id = self.named_root_id(&app.id)?;
                App(NodeApp {
                    id: node_id,
                    args: app.args,
                })
            }
        };
        Ok(ty)
    }

    fn push_fields(&mut self, fields: Vec<syn::Field>) -> Result<Vec<FieldId>, Error> {
        let mut fields = fields;
        let field_ids = fields
            .drain(0..)
            .map(|field| {
                let id = self.fields.len();
                let ty = self.push_type(field.ty)?;
                self.fields.push(Field {
                    name: field.name,
                    ty,
                    constant: field.constant,
                    hidden: field.hidden,
                });
                Ok(FieldId(id))
            })
            .collect::<Result<_, _>>()?;
        Ok(field_ids)
    }
}

impl std::ops::Index<NodeId> for Forest {
    type Output = Node;

    fn index(&self, node_id: NodeId) -> &Self::Output {
        &self.nodes[node_id.0]
    }
}

#[derive(Debug)]
struct BranchIter<'a> {
    forest: &'a Forest,
    node: &'a Node,
    branch_ids: &'a [usize],
}

impl<'a> BranchIter<'a> {
    pub fn is_empty(&self) -> bool {
        self.branch_ids.is_empty()
    }
}

impl<'a> Iterator for BranchIter<'a> {
    type Item = &'a Branch;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some((&branch_id, branch_ids)) = self.branch_ids.split_first() {
            let branch = &self.node.branches[branch_id];
            self.node = &self.forest[branch.target];
            self.branch_ids = branch_ids;
            Some(branch)
        } else {
            None
        }
    }
}

fn resolve_exprs(
    exprs: &[data::Expr],
    fields: &exec::FieldEnv,
) -> Result<Vec<exec::Expr>, ResolveError> {
    exprs
        .iter()
        .map(|expr| resolve_expr(expr, fields))
        .collect::<Result<_, _>>()
}

fn resolve_expr(
    expr: &Expr<R<Ident>>,
    field_env: &exec::FieldEnv,
) -> Result<exec::Expr, ResolveError> {
    use Expr::*;
    let resolve_2 = |args: &[Expr<R<Ident>>; 2]| {
        Ok(Box::new([
            resolve_expr(&args[0], field_env)?,
            resolve_expr(&args[1], field_env)?,
        ]))
    };
    match expr {
        Not(arg) => Ok(Not(Box::new(resolve_expr(arg, field_env)?))),
        And(args) => Ok(And(resolve_exprs(args, field_env)?)),
        Or(args) => Ok(And(resolve_exprs(args, field_env)?)),
        Eq(args) => resolve_2(args).map(Eq),
        Add(args) => resolve_2(args).map(Add),
        Sub(args) => resolve_2(args).map(Sub),
        Mul(args) => resolve_2(args).map(Mul),
        BitNot(arg) => Ok(BitNot(Box::new(resolve_expr(arg, field_env)?))),
        BitAnd(args) => resolve_2(args).map(BitAnd),
        BitOr(args) => resolve_2(args).map(BitOr),
        BitShl(args) => resolve_2(args).map(BitShr),
        BitShr(args) => resolve_2(args).map(BitShl),
        Ref(r) => match r {
            R::Field(name) => {
                let idx = field_env
                    .find_idx(name)
                    .ok_or_else(|| ResolveError::unknown_field(name))?;
                Ok(Ref(R::Field(idx)))
            }
            R::Param(i) => Ok(Ref(R::Param(*i))),
        },
        TypeProperty(type_r, property_name) => match type_r {
            R::Param(i) => Ok(TypeProperty(R::Param(*i), property_name.clone())),
            R::Field(name) => {
                let idx = field_env
                    .find_idx(name)
                    .ok_or_else(|| ResolveError::unknown_field(name))?;
                Ok(TypeProperty(R::Field(idx), property_name.clone()))
            }
        },
        Const(c) => Ok(Const(c.clone())),
    }
}