use ariadne::{ColorGenerator, Fmt, Label, Report, ReportKind, Span as _};
use std::fmt::Debug;

use super::{Span, UnaryOperator};
use crate::typing::{BinaryOperator, Identifier, ModulePath, Type};

#[derive(Debug, PartialEq)]
pub struct TypeError {
    pub file: Option<std::path::PathBuf>,
    pub module: ModulePath,
    pub function: Option<String>,
    pub kind: TypeErrorKind,
}

#[derive(PartialEq, Debug)]
pub struct TypeAndSpan {
    pub ty: Type,
    pub span: Span,
}

#[derive(PartialEq, Debug)]
pub struct BinOpAndSpan {
    pub op: BinaryOperator,
    pub span: Span,
}

#[derive(Debug, PartialEq)]
pub enum TypeErrorKind {
    InvalidBinaryOperator {
        operator: BinOpAndSpan,
        lhs: TypeAndSpan,
        rhs: TypeAndSpan,
    },
    BlockTypeDoesNotMatchFunctionType {
        block_type: Type,
    },
    ReturnTypeDoesNotMatchFunctionType {
        return_expr: Option<TypeAndSpan>,
        return_stmt: TypeAndSpan,
    },
    UnknownIdentifier {
        identifier: String,
    },
    AssignmentMismatch {
        lht: Type,
        rht: Type,
    },
    AssignUndeclared,
    VariableRedeclaration,
    UnknownFunctionCalled(Identifier),
    WrongFunctionArguments,
    ConditionIsNotBool,
    IfElseMismatch,
    InvalidUnaryOperator {
        operator: UnaryOperator,
        inner: Type,
    },
}

impl TypeError {
    pub fn to_report(&self, ast: &crate::ast::Module) -> Report<Span> {
        let mut colors = ColorGenerator::new();
        let c0 = colors.next();
        let c1 = colors.next();
        colors.next();
        let c2 = colors.next();

        match &self.kind {
            TypeErrorKind::InvalidBinaryOperator { operator, lhs, rhs } => {
                Report::build(ReportKind::Error, lhs.span.source, lhs.span.start)
                    .with_message(format!(
                        "Invalid binary operation {} between {} and {}",
                        operator.op.to_string().fg(c0),
                        lhs.ty.to_string().fg(c1),
                        rhs.ty.to_string().fg(c2),
                    ))
                    .with_labels([
                        Label::new(operator.span).with_color(c0),
                        Label::new(lhs.span)
                            .with_message(format!("This has type {}", lhs.ty.to_string().fg(c1)))
                            .with_color(c1)
                            .with_order(2),
                        Label::new(rhs.span)
                            .with_message(format!("This has type {}", rhs.ty.to_string().fg(c2)))
                            .with_color(c2)
                            .with_order(1),
                    ])
                    .finish()
            }

            TypeErrorKind::BlockTypeDoesNotMatchFunctionType { block_type } => {
                let function = ast
                    .functions
                    .iter()
                    .find(|f| f.name == *self.function.as_ref().unwrap())
                    .unwrap();

                let block_color = c0;
                let signature_color = c1;

                let span = function.body.value.as_ref().unwrap().span;

                let func_return_type_text = function
                    .return_type
                    .as_ref()
                    .unwrap_or(&Type::Unit)
                    .to_string()
                    .fg(signature_color);

                let mut report = Report::build(
                    ReportKind::Error,
                    span.source,
                    function.body.value.as_ref().unwrap().span.start,
                )
                .with_message("Function body does not match the signature")
                .with_labels([
                    Label::new(function.body.span.unwrap())
                        .with_message("In this function's body")
                        .with_color(c2),
                    Label::new(span)
                        .with_message(format!(
                            "Returned expression has type {} but the function should return {}",
                            block_type.to_string().fg(block_color),
                            func_return_type_text,
                        ))
                        .with_color(block_color),
                ]);

                if let Some(span) = function.return_type_span {
                    report.add_label(
                        Label::new(span)
                            .with_message(format!("The signature shows {}", func_return_type_text))
                            .with_color(signature_color),
                    );
                }

                report.set_note("The last expression of a function's body is returned");

                if function.return_type.is_none() {
                    report
                        .set_help("You may need to add the return type to the function's signature")
                }

                report.finish()
            }

            TypeErrorKind::ReturnTypeDoesNotMatchFunctionType {
                return_expr,
                return_stmt,
            } => {
                let function = ast
                    .functions
                    .iter()
                    .find(|f| f.name == *self.function.as_ref().unwrap())
                    .unwrap();

                let is_bare_return = return_expr.is_none();

                let mut report = Report::build(
                    ReportKind::Error,
                    *return_stmt.span.source(),
                    return_stmt.span.start,
                )
                .with_message("Return type does not match the function's signature")
                .with_label(
                    Label::new(return_expr.as_ref().unwrap_or(return_stmt).span)
                        .with_color(c1)
                        .with_message(if is_bare_return {
                            format!("Bare return has type {}", Type::Unit.to_string().fg(c1))
                        } else {
                            format!(
                                "This expression has type {}",
                                return_stmt.ty.to_string().fg(c1)
                            )
                        }),
                );

                if let Some(ret_ty_span) = function.return_type_span {
                    report.add_label(Label::new(ret_ty_span).with_color(c0).with_message(format!(
                        "The signature shows {}",
                        function.return_type.as_ref().unwrap().to_string().fg(c0)
                    )))
                }

                if function.return_type.is_none() {
                    report
                        .set_help("You may need to add the return type to the function's signature")
                }

                report.finish()
            }

            _ => todo!(),
        }
    }
}