use nom::{
    branch::alt,
    combinator::map,
    error::context,
    multi::{many0, separated_list0},
    sequence::{delimited, preceded, tuple},
};

use crate::parser::{
    assign_statement, declaration,
    proc::{proc_block, ProcBlock},
    tokens::{token, TokenKind as tk, TokenSpan},
    typename, Assign, IResult, NetDecl, Span,
};

#[derive(Debug)]
pub enum PortDirection {
    Input,
    Output,
}

#[derive(Debug)]
pub struct PortDecl<'a> {
    pub direction: PortDirection,
    pub net: NetDecl<'a>,
}

#[derive(Debug)]
pub struct Module<'a> {
    pub name: Span<'a>,
    pub ports: Vec<PortDecl<'a>>,
    pub items: Vec<ModuleItem<'a>>,
}

#[derive(Debug)]
pub enum ModuleItem<'a> {
    Assign(Assign<'a>),
    Proc(ProcBlock<'a>),
}

fn port_decl(i: TokenSpan) -> IResult<TokenSpan, PortDecl> {
    map(declaration, |net| PortDecl {
        direction: PortDirection::Input,
        net,
    })(i)
}

fn inputs_list(input: TokenSpan) -> IResult<TokenSpan, Vec<PortDecl>> {
    separated_list0(token(tk::Comma), port_decl)(input)
}

fn assign_item(input: TokenSpan) -> IResult<TokenSpan, Assign> {
    context(
        "assignment",
        delimited(token(tk::Assign), assign_statement, token(tk::Semicolon)),
    )(input)
}

fn module_item(input: TokenSpan) -> IResult<TokenSpan, ModuleItem> {
    alt((
        map(assign_item, ModuleItem::Assign),
        map(proc_block, ModuleItem::Proc),
    ))(input)
}

pub fn module(input: TokenSpan) -> IResult<TokenSpan, Module> {
    context(
        "module",
        map(
            tuple((
                token(tk::Module),
                token(tk::Ident),
                delimited(token(tk::LParen), inputs_list, token(tk::RParen)),
                preceded(token(tk::RArrow), typename),
                delimited(token(tk::LBrace), many0(module_item), token(tk::RBrace)),
            )),
            |(_, name, inputs, _ret, items)| Module {
                // TODO: bring back returns
                name: name.span(),
                ports: inputs,
                items,
            },
        ),
    )(input)
}

#[cfg(test)]
mod test {
    use super::*;
    use nom::combinator::all_consuming;

    #[test]
    fn test_decl() {
        declaration("reg abcd".into()).unwrap();
    }

    #[test]
    fn test_assignment_item() {
        all_consuming(assign_item)(" assign a = b ; ".into()).unwrap();
        all_consuming(assign_item)(" assign a = b | c ; ".into()).unwrap();
    }

    #[test]
    fn test_module_item() {
        all_consuming(module_item)(" assign a = b ;".into()).unwrap();
    }
}