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

// 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,
            (Type::Nothing, Type::Nothing) => 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 Declaration {
    pub declaration_type: Type,
}

pub struct Environment {
    pub parent: Option<EnvironmentRef>,
    pub declarations: HashMap<Box<str>, Declaration>,
}

impl Environment {
    pub fn new(parent: Option<EnvironmentRef>) -> Self {
        Environment {
            parent,
            declarations: HashMap::new(),
        }
    }

    pub fn bind(&self, token: &Token) -> Option<&Declaration> {
        if let Some(decl) = self.declarations.get(token.as_str()) {
            return Some(decl);
        }

        let mut current = &self.parent;
        while let Some(env) = current {
            if let Some(decl) = env.declarations.get(token.as_str()) {
                return Some(decl);
            }
            current = &env.parent;
        }

        None
    }
}

#[derive(Clone)]
pub struct EnvironmentRef(Rc<Environment>);

impl EnvironmentRef {
    pub fn new(environment: Environment) -> Self {
        EnvironmentRef(Rc::new(environment))
    }
}

impl std::ops::Deref for EnvironmentRef {
    type Target = Environment;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

fn set_logical_parents(
    parents: &mut Vec<Option<TreeRef>>,
    syntax_tree: &SyntaxTree,
    t: TreeRef,
    parent: Option<TreeRef>,
) {
    parents[t.index()] = parent.clone();

    let tree = &syntax_tree[t];
    // eprintln!("SET PARENT {parent:?} => CHILD {tree:?} ({t:?})");
    match tree.kind {
        TreeKind::Block | TreeKind::File => {
            // In a block (or at the top level), each child actually points
            // to the previous child as the logical parent, so that variable
            // declarations that occur as part of statements in the block are
            // available to statements later in the block.
            let mut parent = Some(t);
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => {
                        set_logical_parents(parents, syntax_tree, *ct, parent);
                        parent = Some(*ct);
                    }
                }
            }
        }
        TreeKind::LetStatement => {
            // In a let statement, the logical parent of the children is
            // actually the logical parent of the let statement, so that the
            // variable doesn't have itself in scope. :P
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => set_logical_parents(parents, syntax_tree, *ct, parent),
                }
            }
        }
        _ => {
            // By default, the parent for each child is current tree.
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => set_logical_parents(parents, syntax_tree, *ct, Some(t)),
                }
            }
        }
    }
}

enum Incremental<T> {
    InProgress,
    Complete(T),
}

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,

    // Instead of physical parents, this is the set of *logical* parents.
    // This is what is used for binding.
    logical_parents: Vec<Option<TreeRef>>,

    // TODO: State should be externalized instead of this refcell nonsense.
    errors: RefCell<Vec<Error>>,
    types: RefCell<HashMap<TreeRef, Incremental<Type>>>,
    environments: RefCell<HashMap<TreeRef, Incremental<EnvironmentRef>>>,
    empty_environment: EnvironmentRef,
}

impl<'a> Semantics<'a> {
    pub fn new(tree: &'a SyntaxTree<'a>, lines: &'a Lines) -> Self {
        let mut logical_parents = Vec::with_capacity(tree.len());
        logical_parents.resize(tree.len(), None);
        if let Some(root) = tree.root() {
            set_logical_parents(&mut logical_parents, tree, root, None);
        }

        let mut semantics = Semantics {
            syntax_tree: tree,
            lines,
            logical_parents,
            errors: RefCell::new(vec![]),
            types: RefCell::new(HashMap::new()),
            environments: RefCell::new(HashMap::new()),
            empty_environment: EnvironmentRef::new(Environment::new(None)),
        };

        // 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()
    }

    pub fn logical_parent(&self, tr: TreeRef) -> Option<TreeRef> {
        if tr.index() < self.logical_parents.len() {
            self.logical_parents[tr.index()]
        } else {
            None
        }
    }

    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 environment_of(&self, t: TreeRef) -> EnvironmentRef {
        match self.environments.borrow().get(&t) {
            None => (),
            Some(Incremental::Complete(e)) => return e.clone(),
            Some(Incremental::InProgress) => {
                // TODO: Rewrite as complete with empty after reporting error.
                // eprintln!("environment_of circular => {t:?}");
                self.report_error_tree_ref(
                    t,
                    "INTERNAL COMPILER ERROR: Circular dependency detected: environment",
                );
                return self.empty_environment.clone();
            }
        }
        self.environments
            .borrow_mut()
            .insert(t, Incremental::InProgress);

        let tree = &self.syntax_tree[t];
        // eprintln!("environment_of => {tree:?}");

        let parent = match self.logical_parents[t.index()] {
            Some(t) => self.environment_of(t),
            None => self.empty_environment.clone(),
        };

        let result = match tree.kind {
            TreeKind::LetStatement => self.environment_of_let(parent, tree),
            TreeKind::FunctionDecl => self.environment_of_func(parent, tree),

            // TODO: MORE Things that introduce an environment!
            _ => parent,
        };

        self.environments
            .borrow_mut()
            .insert(t, Incremental::Complete(result.clone()));
        result
    }

    fn environment_of_let(&self, parent: EnvironmentRef, tree: &Tree) -> EnvironmentRef {
        let Some(name) = tree.nth_token(1) else {
            return parent; // Error is already reported?
        };

        let declaration_type = match tree.nth_tree(3) {
            Some(expr) => self
                .type_of(expr)
                .expect("the tree in the expression should yield a type"),

            // The syntax error should already have been reported, so we'll
            // stick with error type here. (But bind the name, because we see
            // it!)
            None => Type::Error,
        };

        let mut environment = Environment::new(Some(parent));
        environment
            .declarations
            .insert(name.as_str().into(), Declaration { declaration_type });

        EnvironmentRef::new(environment)
    }

    fn environment_of_func(&self, parent: EnvironmentRef, tree: &Tree) -> EnvironmentRef {
        let Some(param_list) = tree.nth_tree(2) else {
            return parent; // SE
        };
        let param_list = &self.syntax_tree[param_list];
        if param_list.kind != TreeKind::ParamList {
            return parent; // SE
        }

        let mut environment = Environment::new(Some(parent));
        for child in param_list.children.iter() {
            let Child::Tree(ct) = child else {
                continue;
            };
            let param = &self.syntax_tree[*ct];
            if param.kind != TreeKind::Parameter {
                continue;
            }

            let Some(param_name) = param.nth_token(0) else {
                continue;
            };

            let declaration_type = if let Some(type_expression) = param.nth_tree(2) {
                self.type_of(type_expression)
                    .expect("the type expression should yield *some* type here")
            } else {
                Type::Error
            };

            environment
                .declarations
                .insert(param_name.as_str().into(), Declaration { declaration_type });
        }

        EnvironmentRef::new(environment)
    }

    pub fn type_of(&self, t: TreeRef) -> Option<Type> {
        match self.types.borrow().get(&t) {
            None => (),
            Some(Incremental::Complete(existing)) => return Some(existing.clone()),
            Some(Incremental::InProgress) => {
                // TODO: Rewrite as complete with error after reporting error.
                // eprintln!("type_of circular => {t:?}");
                self.report_error_tree_ref(
                    t,
                    "INTERNAL COMPILER ERROR: Circular dependency detected: type",
                );
                return Some(Type::Error);
            }
        }
        self.types.borrow_mut().insert(t, Incremental::InProgress);

        let tree = &self.syntax_tree[t];
        // eprintln!("type_of => {tree:?}");

        let result = match tree.kind {
            TreeKind::Error => Some(Type::Error),
            TreeKind::UnaryExpression => self.type_of_unary(tree),
            TreeKind::BinaryExpression => self.type_of_binary(tree),
            TreeKind::TypeExpression => self.type_of_type_expr(tree),
            TreeKind::Block => self.type_of_block(tree),
            TreeKind::LiteralExpression => self.type_of_literal(tree),
            TreeKind::GroupingExpression => self.type_of_grouping(tree),
            TreeKind::ConditionalExpression => self.type_of_conditional(tree),
            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),
            TreeKind::Identifier => self.type_of_identifier(tree),

            // TODO: Previously I had short-circuited here and not put anything
            //       in the table if this node isn't the kind that I would
            //       normally compute a type for. I should keep doing that to
            //       detect nonsense without blowing out the hash table. If
            //       we're going to be computing a type for every node it
            //       should just be an array instead of a hash table.
            _ => None,
        };

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

        self.types
            .borrow_mut()
            .insert(t, Incremental::Complete(result.clone()));
        Some(result)
    }

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

        // TODO: This will *clearly* need to get better.
        let token = tree.nth_token(0)?;
        match token.as_str() {
            "f64" => Some(Type::F64),
            "string" => Some(Type::String),
            "bool" => Some(Type::Bool),
            "()" => Some(Type::Nothing),
            _ => Some(Type::Error),
        }
    }

    fn type_of_block(&self, tree: &Tree) -> 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 {
            // TODO: if `is_unreachable` here then we actually have
            // unreachable code here! We should warn about it I guess.

            is_unreachable = self
                .type_of(tree.nth_tree(i)?)
                .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)?)
            .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) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::GroupingExpression);

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

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

        let cond_tree = tree.nth_tree(1)?;
        let cond_type = self.type_of(cond_tree).expect("must be expression");
        let then_type = self.type_of(tree.nth_tree(2)?).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)?).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),

                (Type::Unreachable, None) => Some(Type::Unreachable),
                (Type::Unreachable, Some(t)) => Some(t),
                (t, Some(Type::Unreachable)) => Some(t),

                (then_type, else_type) => {
                    let else_type = else_type.unwrap_or(Type::Nothing);
                    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)
                    }
                }
            }
        }
    }

    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) -> 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)?).expect("must be expression");
        Some(match expression_type {
            Type::Unreachable => Type::Unreachable,
            _ => {
                // A semicolon at the end of an expression statement discards
                // the value, leaving us with nothing. (Even if the
                // expression otherwise generated a type error!)
                if last_is_semicolon {
                    Type::Nothing
                } else {
                    expression_type
                }
            }
        })
    }

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

        let id = tree.nth_token(0)?;
        if let Some(parent) = tree.parent {
            let environment = self.environment_of(parent);
            if let Some(declaration) = environment.bind(id) {
                return Some(declaration.declaration_type);
            }
        }

        self.report_error_tree(tree, format!("cannot find value {id} here"));
        Some(Type::Error)
    }
}