use celestial_hub_astrolabe::{
  ast::{
    DataSection, Instruction, InstructionArgument, Register, Statement, TextSection, Variable,
  },
  lexer::tokens::{Type, Value},
};

use crate::ast::Statement as CompassStatement;

use super::Codegen;

pub struct MipsCodegen;

impl Codegen for MipsCodegen {
  fn generate(&self, ast: Vec<CompassStatement>) -> Result<String, String> {
    let mut data_section = celestial_hub_astrolabe::ast::DataSection::default();
    let mut text_section = celestial_hub_astrolabe::ast::TextSection {
      entrypoint: "main".to_string(),
      ..Default::default()
    };

    let mut register_counter = 0;
    let mut register_map = std::collections::HashMap::new();

    for statement in ast {
      match statement {
        CompassStatement::VariableDeclaration(var) => {
          let register = format!("${}", register_counter);
          register_counter += 1;
          register_map.insert(var.name.clone(), register.clone());

          match var.value {
            crate::ast::Expr::Operand(op) => match op {
              crate::ast::Operand::LiteralI8(val) => {
                load_immediate(&mut text_section, register, val as u32)
              }
              crate::ast::Operand::LiteralI16(val) => {
                load_immediate(&mut text_section, register, val as u32)
              }
              crate::ast::Operand::LiteralI32(val) => {
                load_immediate(&mut text_section, register, val as u32)
              }
              crate::ast::Operand::LiteralI64(val) => {
                load_immediate(&mut text_section, register, val as u32)
              }
              crate::ast::Operand::LiteralU8(val) => {
                load_immediate(&mut text_section, register, val as u32)
              }
              crate::ast::Operand::LiteralU16(val) => {
                load_immediate(&mut text_section, register, val as u32)
              }
              crate::ast::Operand::LiteralU32(val) => {
                load_immediate(&mut text_section, register, val)
              }
              crate::ast::Operand::LiteralU64(val) => {
                return Err("Cannot store a 64-bit integer in a 32-bit register".to_string());
              }
              crate::ast::Operand::LiteralStr(val) => {
                load_string(&mut text_section, &mut data_section, register, val)
              }
              crate::ast::Operand::LiteralBool(val) => todo!(),
              crate::ast::Operand::LiteralF32(val) => todo!(),
              crate::ast::Operand::LiteralF64(val) => todo!(),
              crate::ast::Operand::Identifier(var) => todo!(),
            },
            crate::ast::Expr::BinaryOperation(bin_op) => match bin_op {
              crate::ast::BinaryOperation::Arithmetic {
                lhs, operator, rhs, ..
              } => {
                if let (
                  crate::ast::Operand::Identifier(lhs),
                  crate::ast::Operand::Identifier(rhs),
                ) = (lhs, rhs)
                {
                  let lhs_register = register_map.get(&lhs).unwrap().clone();
                  let rhs_register = register_map.get(&rhs).unwrap().clone();

                  // Depending on the operator, we need to use different instructions
                  text_section.statements.push(
                    celestial_hub_astrolabe::ast::Statement::Instruction(match operator {
                      crate::ast::Operator::Add => celestial_hub_astrolabe::ast::Instruction::Add(
                        [
                          celestial_hub_astrolabe::ast::InstructionArgument::Register(
                            celestial_hub_astrolabe::ast::Register { name: register },
                          ),
                          celestial_hub_astrolabe::ast::InstructionArgument::Register(
                            celestial_hub_astrolabe::ast::Register { name: lhs_register },
                          ),
                          celestial_hub_astrolabe::ast::InstructionArgument::Register(
                            celestial_hub_astrolabe::ast::Register { name: rhs_register },
                          ),
                        ]
                        .into(),
                      ),
                      crate::ast::Operator::Sub => todo!(),
                      crate::ast::Operator::Mul => todo!(),
                      crate::ast::Operator::Div => todo!(),
                    }),
                  );
                }
              }
              crate::ast::BinaryOperation::Conditional {
                lhs,
                condition,
                rhs,
                operation_type,
              } => todo!(),
            },
            crate::ast::Expr::FunctionCall(_) => todo!(),
          }
        }
        CompassStatement::ConditionalJump {
          condition,
          label,
          location,
        } => todo!(),
        CompassStatement::UnconditionalJump { label, location } => todo!(),
        CompassStatement::Label { name, location } => todo!(),
        CompassStatement::FunctionDefinition(_) => todo!(),
      }
    }

    let program = celestial_hub_astrolabe::ast::Program {
      data_section,
      text_section,
    };

    Ok(program.to_string())
  }
}

fn load_immediate(text_section: &mut TextSection, register: String, value: u32) {
  text_section
    .statements
    .push(Statement::Instruction(Instruction::Li(
      [
        InstructionArgument::Register(Register { name: register }),
        InstructionArgument::Immediate(value),
      ]
      .into(),
    )));
}

fn load_string(
  text_section: &mut TextSection,
  data_section: &mut DataSection,
  register: String,
  value: String,
) {
  // Search for an existing string label
  let mut label = String::new();
  let mut found = false;
  for (i, variable) in data_section.variables.iter().enumerate() {
    if let Value::String(existing_value) = &variable.value {
      if existing_value == &value {
        label = format!("str_{}", i);
        found = true;
        break;
      }
    }
  }

  // If the string is not found, add it to the data section
  if !found {
    label = format!("str_{}", data_section.variables.len());
    data_section.variables.push(Variable {
      name: label.clone(),
      type_: Type::Asciiz,
      value: Value::String(value),
    });
  }

  // Load the address of the string into the register
  text_section
    .statements
    .push(Statement::Instruction(Instruction::La(
      [
        InstructionArgument::Register(Register { name: register }),
        InstructionArgument::Label(label),
      ]
      .into(),
    )));
}