use core::fmt;
use std::process;

use ariadne::{Color, Label, Report, ReportKind, Source};
use chumsky::{
    input::{Stream, ValueInput},
    prelude::*,
};
use logos::Logos;

use crate::inst::Inst;

#[derive(Logos, Debug, Clone, PartialEq, Eq, Hash)]
#[logos(skip r"[ \t\n\r\f]+")]
enum Token<'a> {
    #[token("->")]
    #[token("→")]
    Set,

    #[token("(")]
    OpenParen,

    #[token(")")]
    CloseParen,

    #[regex(r"[a-zA-Z0-9+*/%.:=-]+")]
    VarOrNum(&'a str),

    #[token("!")]
    Call,

    Error,
}

impl fmt::Display for Token<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Token::Set => write!(f, "→"),
            Token::OpenParen => write!(f, "("),
            Token::CloseParen => write!(f, ")"),
            Token::VarOrNum(s) => write!(f, "{s}"),
            Token::Call => write!(f, "!"),
            Token::Error => write!(f, "ERROR"),
        }
    }
}

fn parser<'tokens, 'src: 'tokens, I>()
-> impl Parser<'tokens, I, Vec<Inst<'src>>, extra::Err<Rich<'tokens, Token<'src>>>>
where
    I: ValueInput<'tokens, Token = Token<'src>, Span = SimpleSpan>,
{
    let inst = recursive(|inst| {
        let block = inst
            .clone()
            .repeated()
            .collect()
            .map(Inst::Block)
            .delimited_by(
                just(Token::OpenParen),
                just(Token::CloseParen).or(end().validate(|(), e, emitter| {
                    emitter.emit(Rich::custom(e.span(), "Expected close paren"));
                    Token::Error
                })),
            );

        choice((
            just(Token::Set).ignore_then(inst.validate(|v: Inst<'_>, e, emitter| {
                if let Inst::Var(v) = v {
                    Inst::Set(v)
                } else {
                    emitter.emit(Rich::custom(
                        e.span(),
                        "Set target must be a variable".to_string(),
                    ));
                    Inst::Set("invalid")
                }
            })),
            just(Token::Call).to(Inst::Call),
            select! {
                Token::VarOrNum(s) => s.parse::<f64>().map_or_else(|_| Inst::Var(s), Inst::Num),
            },
            block,
            just(Token::Error).validate(|_, e, emitter| {
                emitter.emit(Rich::custom(e.span(), "Invalid token".to_string()));
                Inst::Call
            }),
        ))
    });

    inst.or(
        just(Token::CloseParen).try_map(|_, span| Err(Rich::custom(span, "Unmatched close paren")))
    )
    .repeated()
    .collect()
}

pub fn parse(src: &str) -> Vec<Inst<'_>> {
    let token_iter = Token::lexer(src).spanned().map(|(tok, span)| match tok {
        Ok(tok) => (tok, span.into()),
        Err(()) => (Token::Error, span.into()),
    });

    let token_stream =
        Stream::from_iter(token_iter).map((0..src.len()).into(), |(t, s): (_, _)| (t, s));

    let source_filename = "input";

    match parser().parse(token_stream).into_result() {
        Ok(insts) => insts,
        Err(errs) => {
            for err in errs {
                Report::build(
                    ReportKind::Error,
                    (source_filename, err.span().into_range()),
                )
                .with_config(ariadne::Config::new().with_index_type(ariadne::IndexType::Byte))
                .with_label(
                    Label::new((source_filename, err.span().into_range()))
                        .with_message(err.reason())
                        .with_color(Color::Red),
                )
                .finish()
                .eprint((source_filename, Source::from(&src)))
                .unwrap();
            }
            process::exit(2);
        }
    }
}