use crate::{
    parser::{Child, SyntaxTree, Tree, TreeKind, TreeRef},
    tokens::{Lines, TokenKind},
};
use std::{cell::RefCell, collections::HashMap, fmt};

// TODO: An error should have:
//
//   - a start
//   - an end
//   - a focus
//   - descriptive messages
//
// that will have to wait for now
#[derive(Clone, PartialEq, Eq)]
pub struct Error {
    pub start: (usize, usize),
    pub end: (usize, usize),
    pub message: String,
}

impl Error {
    pub fn new<T>(line: usize, column: usize, message: T) -> Self
    where
        T: ToString,
    {
        Error {
            start: (line, column),
            end: (line, column),
            message: message.to_string(),
        }
    }

    pub fn new_spanned<T>(start: (usize, usize), end: (usize, usize), message: T) -> Self
    where
        T: ToString,
    {
        Error {
            start,
            end,
            message: message.to_string(),
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self}")
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}: {}", self.start.0, self.start.1, self.message)
    }
}

#[derive(Copy, Clone)]
pub enum Type {
    // Signals a type error. If you receive this then you know that an error
    // has already been reported; if you produce this be sure to also note
    // the error in the errors collection.
    Error,

    // Signals that the expression has a control-flow side-effect and that no
    // value will ever result from this expression. Usually this means
    // everything's fine.
    Unreachable,

    Nothing,
    // TODO: Numeric literals should be implicitly convertable, unlike other
    //       types. Maybe just "numeric literal" type?
    F64,
    String,
    Bool,
}

impl Type {
    pub fn is_error(&self) -> bool {
        match self {
            Type::Error => true,
            _ => false,
        }
    }

    pub fn compatible_with(&self, other: &Type) -> bool {
        // TODO: This is wrong; we because of numeric literals etc.
        match (self, other) {
            (Type::F64, Type::F64) => true,
            (Type::String, Type::String) => true,
            (Type::Bool, Type::Bool) => true,
            (Type::Unreachable, Type::Unreachable) => true,

            // Avoid introducing more errors
            (Type::Error, _) => true,
            (_, Type::Error) => true,

            (_, _) => false,
        }
    }
}

impl fmt::Debug for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self}")
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use Type::*;
        match self {
            Error => write!(f, "<< INTERNAL ERROR >>"),
            Unreachable => write!(f, "<< UNREACHABLE >>"),
            Nothing => write!(f, "()"),
            F64 => write!(f, "f64"),
            String => write!(f, "string"),
            Bool => write!(f, "bool"),
        }
    }
}

pub struct Semantics<'a> {
    // TODO: Do I really want my own copy here? Should we standardize on Arc
    //       or Rc or some other nice sharing mechanism?
    syntax_tree: &'a SyntaxTree<'a>,
    lines: &'a Lines,
    errors: RefCell<Vec<Error>>,
    types: RefCell<HashMap<(TreeRef, bool), Type>>,
}

impl<'a> Semantics<'a> {
    pub fn new(tree: &'a SyntaxTree<'a>, lines: &'a Lines) -> Self {
        let mut semantics = Semantics {
            syntax_tree: tree,
            lines,
            errors: RefCell::new(vec![]),
            types: RefCell::new(HashMap::new()),
        };

        // NOTE: We ensure all the known errors are reported before we move
        //       on to answering any other questions. We're going to work as
        //       hard as we can from a partial tree.
        if let Some(tr) = semantics.syntax_tree.root() {
            semantics.gather_errors(tr);
        }

        semantics
    }

    pub fn tree(&self) -> &SyntaxTree<'a> {
        &self.syntax_tree
    }

    pub fn snapshot_errors(&self) -> Vec<Error> {
        (*self.errors.borrow()).clone()
    }

    fn report_error<T>(&self, position: usize, error: T)
    where
        T: ToString,
    {
        let (line, col) = self.lines.position(position);
        self.errors
            .borrow_mut()
            .push(Error::new(line, col, error.to_string()));
    }

    fn report_error_span<T>(&self, start: usize, end: usize, error: T)
    where
        T: ToString,
    {
        let start = self.lines.position(start);
        let end = self.lines.position(end);
        self.errors
            .borrow_mut()
            .push(Error::new_spanned(start, end, error.to_string()));
    }

    fn report_error_tree<T>(&self, tree: &Tree<'a>, error: T)
    where
        T: ToString,
    {
        self.report_error_span(tree.start_pos, tree.end_pos, error)
    }

    fn report_error_tree_ref<T>(&self, tree: TreeRef, error: T)
    where
        T: ToString,
    {
        let tree = &self.syntax_tree[tree];
        self.report_error_span(tree.start_pos, tree.end_pos, error)
    }

    fn gather_errors(&mut self, tree: TreeRef) {
        let mut stack = vec![tree];
        while let Some(tr) = stack.pop() {
            let tree = &self.syntax_tree[tr];
            for child in &tree.children {
                match child {
                    Child::Token(t) => {
                        if t.kind == TokenKind::Error {
                            self.report_error(t.start, t.as_str());
                        }
                    }
                    Child::Tree(t) => stack.push(*t),
                }
            }
        }
    }

    pub fn type_of(&self, t: TreeRef, value_required: bool) -> Option<Type> {
        if let Some(existing) = self.types.borrow().get(&(t, value_required)) {
            return Some(existing.clone());
        }

        let tree = &self.syntax_tree[t];
        let result = match tree.kind {
            TreeKind::Error => Some(Type::Error),
            TreeKind::UnaryExpression => self.type_of_unary(tree, value_required),
            TreeKind::BinaryExpression => self.type_of_binary(tree, value_required),
            TreeKind::TypeExpression => self.type_of_type_expr(tree, value_required),
            TreeKind::Block => self.type_of_block(tree, value_required),
            TreeKind::LiteralExpression => self.type_of_literal(tree),
            TreeKind::GroupingExpression => self.type_of_grouping(tree, value_required),
            TreeKind::ConditionalExpression => self.type_of_conditional(tree, value_required),
            TreeKind::CallExpression => self.type_of_call(tree),
            TreeKind::Argument => self.type_of_argument(tree),

            TreeKind::LetStatement => Some(Type::Nothing),
            TreeKind::ReturnStatement => Some(Type::Unreachable),
            TreeKind::ExpressionStatement => {
                self.type_of_expression_statement(tree, value_required)
            }
            TreeKind::Identifier => self.type_of_identifier(tree),
            _ => return None,
        };

        // NOTE: These return `None` if they encounter some problem.
        let result = result.unwrap_or(Type::Error);

        self.types
            .borrow_mut()
            .insert((t, value_required), result.clone());
        Some(result)
    }

    fn type_of_unary(&self, tree: &Tree, value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::UnaryExpression);

        let op = tree.nth_token(0)?;
        let expr = tree.nth_tree(1)?;

        let argument_type = self
            .type_of(expr, value_required)
            .expect("Our argument should be an expression");

        match (op.kind, argument_type) {
            (TokenKind::Plus, Type::F64) => Some(Type::F64),
            (TokenKind::Minus, Type::F64) => Some(Type::F64),
            (TokenKind::Bang, Type::Bool) => Some(Type::Bool),

            // This is dumb and should be punished, probably.
            (_, Type::Unreachable) => {
                self.report_error(
                    op.start,
                    "cannot apply a unary operator to something that doesn't yield a value",
                );
                Some(Type::Error)
            }

            // Propagate existing errors without additional complaint.
            (_, Type::Error) => Some(Type::Error),

            (_, arg_type) => {
                self.report_error(
                    op.start,
                    format!(
                        "cannot apply unary operator '{}' to value of type {}",
                        op.as_str(),
                        arg_type
                    ),
                );
                Some(Type::Error)
            }
        }
    }

    fn type_of_binary(&self, tree: &Tree, value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::BinaryExpression);
        let lhs = self
            .type_of(tree.nth_tree(0)?, value_required)
            .expect("must be an expression");
        let op = tree.nth_token(1)?;
        let rhs = self
            .type_of(tree.nth_tree(2)?, value_required)
            .expect("must be an expression");

        match (op.kind, lhs, rhs) {
            (
                TokenKind::Plus | TokenKind::Minus | TokenKind::Star | TokenKind::Slash,
                Type::F64,
                Type::F64,
            ) => Some(Type::F64),

            (TokenKind::Plus, Type::String, Type::String) => Some(Type::String),

            (TokenKind::And | TokenKind::Or, Type::Bool, Type::Bool) => Some(Type::Bool),

            // This is dumb and should be punished, probably.
            (_, _, Type::Unreachable) => {
                self.report_error(
                    op.start,
                    format!("cannot apply '{op}' to an argument that doesn't yield a value (on the right)"),
                );
                Some(Type::Error)
            }
            (_, Type::Unreachable, _) => {
                self.report_error(
                    op.start,
                    format!("cannot apply '{op}' to an argument that doesn't yield a value (on the left)"),
                );
                Some(Type::Error)
            }

            // Propagate existing errors without additional complaint.
            (_, Type::Error, _) => Some(Type::Error),
            (_, _, Type::Error) => Some(Type::Error),

            // Missed the whole table, it must be an error.
            (_, left_type, right_type) => {
                self.report_error(
                    op.start,
                    format!("cannot apply binary operator '{op}' to expressions of type '{left_type}' (on the left) and '{right_type}' (on the right)"),
                );
                Some(Type::Error)
            }
        }
    }

    fn type_of_type_expr(&self, tree: &Tree, _value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::TypeExpression);
        Some(Type::Error)
    }

    fn type_of_block(&self, tree: &Tree, value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::Block);

        if tree.children.len() < 2 {
            return None;
        }

        // if tree.children.len() == 2 {
        //     // Empty blocks generate Nothing.
        //     return Some(Type::Nothing);
        // }

        // The type of the block is the type of the last expression.
        // (But the last child is the closing brace probably?)
        let last_is_brace = tree.nth_token(tree.children.len() - 1).is_some();
        let last_index = tree.children.len() - if last_is_brace { 2 } else { 1 };

        let mut is_unreachable = false;
        for i in 1..last_index {
            is_unreachable = self
                .type_of(tree.nth_tree(i)?, false)
                .map(|t| matches!(t, Type::Unreachable))
                .unwrap_or(false)
                || is_unreachable;
        }

        // NOTE: If for some reason the last statement is unsuitable for a
        //       type then we consider the type of the block to be Nothing.
        //       (And explicitly not Error, which is what returning None
        //       would yield.)
        let last_type = self
            .type_of(tree.nth_tree(last_index)?, value_required)
            .unwrap_or(Type::Nothing);

        // If anything in this block generated an "Unreachable" then the
        // whole type of the block is "unreachable" no matter what.
        Some(if is_unreachable {
            Type::Unreachable
        } else {
            last_type
        })
    }

    fn type_of_literal(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::LiteralExpression);

        let tok = tree.nth_token(0)?;
        let pig = match tok.kind {
            TokenKind::Number => Type::F64,
            TokenKind::String => Type::String,
            TokenKind::True | TokenKind::False => Type::Bool,
            _ => panic!("the token {tok} doesn't have a type!"),
        };
        Some(pig)
    }

    fn type_of_grouping(&self, tree: &Tree, value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::GroupingExpression);

        let expr = tree.nth_tree(1)?;
        Some(
            self.type_of(expr, value_required)
                .expect("the thing in the parenthesis must have some type"),
        )
    }

    fn type_of_conditional(&self, tree: &Tree, value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::ConditionalExpression);

        let cond_tree = tree.nth_tree(1)?;
        let cond_type = self.type_of(cond_tree, true).expect("must be expression");
        let then_type = self
            .type_of(tree.nth_tree(2)?, value_required)
            .expect("must be expression");

        let has_else = tree
            .nth_token(3)
            .map(|t| t.kind == TokenKind::Else)
            .unwrap_or(false);
        let else_type = if has_else {
            Some(
                self.type_of(tree.nth_tree(4)?, value_required)
                    .expect("must be expression"),
            )
        } else {
            None
        };

        if !cond_type.compatible_with(&Type::Bool) {
            if !cond_type.is_error() {
                self.report_error_tree_ref(cond_tree, "conditions must yield a boolean");
            }
            Some(Type::Error)
        } else {
            match (then_type, else_type) {
                (Type::Error, _) => Some(Type::Error),
                (_, Some(Type::Error)) => Some(Type::Error),
                (_, None) if value_required => {
                    self.report_error_tree(
                        tree,
                        "this conditional expression needs an else arm to produce a value",
                    );
                    Some(Type::Error)
                }
                (then_type, Some(else_type)) if value_required => {
                    if else_type.compatible_with(&Type::Unreachable) {
                        // Doesn't matter if the value is required; the else branch
                        // will never generate a value for us so let's ignore it.
                        Some(then_type)
                    } else if then_type.compatible_with(&Type::Unreachable) {
                        // Or the then branch is unreachable, same thing with else
                        // then.
                        Some(else_type)
                    } else if !then_type.compatible_with(&else_type) {
                        self.report_error_tree(
                            tree,
                            format!("the type of the `then` branch ({then_type}) must match the type of the `else` branch ({else_type})"),
                        );
                        Some(Type::Error)
                    } else {
                        Some(then_type)
                    }
                }
                (_, _) => {
                    assert!(!value_required);
                    Some(Type::Unreachable)
                }
            }
        }
    }

    fn type_of_call(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::CallExpression);
        Some(Type::Error)
    }

    fn type_of_argument(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::Argument);
        Some(Type::Error)
    }

    fn type_of_expression_statement(&self, tree: &Tree, value_required: bool) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::ExpressionStatement);
        let last_is_semicolon = tree
            .nth_token(tree.children.len() - 1)
            .map(|t| t.kind == TokenKind::Semicolon)
            .unwrap_or(false);

        let expression_type = self
            .type_of(tree.nth_tree(0)?, value_required && !last_is_semicolon)
            .expect("must be expression");
        Some(match expression_type {
            Type::Error => Type::Error,
            Type::Unreachable => Type::Unreachable,
            _ => {
                // A semicolon at the end of an expression statement discards
                // the value, leaving us with nothing.
                if last_is_semicolon {
                    Type::Nothing
                } else {
                    expression_type
                }
            }
        })
    }

    fn type_of_identifier(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::Identifier);
        Some(Type::Error)
    }
}