use num_traits::pow::Pow;
use std::collections::{ BTreeMap, HashMap };
// use std::env; // TODO: Get args from here
use std::io::{ self, BufRead };
use std::fmt;
use std::rc::Rc;
use rust_decimal::prelude::*;
use rust_decimal_macros::dec;

mod tests;


const HELP_TEXT: &str =
    "LITERALS\n\
     0-9     integers\n\
     _       toggle sign (or, as a command, negate the top item on the stack)\n\
     [...]   string\n\
     (...)   named command (see below)\n\
     \n\
     ARITHMETIC OPERATIONS\n\
     +       add the top two items on the stack\n\
     -       subtract the top item on the stack from the second item\n\
     *       multiply the top two items on the stack\n\
     /       divide the second item on the stack by the top item\n\
     %       take the modulus of the second item by the base of the top item\n\
     ^       raise the second item to the power of the top item\n\
     v       take the square root of the top item on the stack\n\
     \n\
     STACK OPERATIONS\n\
     d       duplicate the top item on the stack\n\
     c       clear the stack\n\
     ,       drop the top item from the stack\n\
     !       assert the stack height, with the top item as desired height\n\
     \n\
     REGISTER OPERATIONS\n\
     l       load from register (register name as next character)\n\
     s       store to register (register name as next character)\n\
     \n\
     CONTROL FLOW\n\
     x       evaluate string\n\
     \n\
     META\n\
     p       print the top item of the stack, leaving it in place\n\
     n       print the top item of the stack, popping it\n\
     q       quit (you can also ^D)\n\
     ?       help (this message)\n\
     \n\
     Good luck!\n\n";


#[derive(Debug,PartialEq,Eq,Clone)]
enum Value {
  Str(String),
  Num(Decimal),
}

impl fmt::Display for Value {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    match self {
      Value::Str(s) => write!(f, "[{}]", s),
      Value::Num(d) => write!(f, "{}", d),
    }
  }
}

fn print_value(v: &Value) {
  match v {
    Value::Str(s) => print!("{}", s),
    Value::Num(d) => print!("{}", d),
  }
}


#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Mode {
  Integer,
  Decimal,
  Str,
  CommandChar,
  CommandNamed,
  RegisterChar,
  RegisterStr,
}


#[derive(Debug)]
enum Exit {
  WithMessage(String),
  Quit,
}

fn err_msg(msg: String) -> Exit {
  Exit::WithMessage(msg)
}


fn baseline_registers() -> HashMap<String, Value> {
  let kvs = [
    ("F->C", "1! 32- 5 9/*"),
    ("C->F", "1! 9 5/* 32+"),
    ("C->K", "1! 273+"),
    ("K->C", "1! 273-"),
    ("Km->mi", "1! 1.609344/"),
    ("mi->Km", "1! 1.609344*"),
  ].map(|kv: (&str, &str)| (String::from(kv.0), Value::Str(kv.1.into())));

  HashMap::from(kvs)
}


fn main() {
  let mut state = RPState {
    registers: baseline_registers(),
    stack: Vec::new(),
    mode: Mode::CommandChar,
    readtables: make_default_readtables(),
    wip_str: String::from(""),
    // TODO: I don't think we need a return stack
    // But this the closest thing we have right now
    reg_command: String::from(""),
    eat_count: 0,
    num: dec!(0.0),
    is_num_negative: false,
    decimal_offset: dec!(1),
  };
  let stdin = io::stdin();
  for line in stdin.lock().lines() {
    match eval(&line.unwrap(), &mut state) {
      Ok(()) => {
        let mut first = true;
        print!("{}", "stack: {");
        for elem in state.stack.iter() {
          if !first { print!("{}", ", "); }
          print!("{}", elem);
          first = false;
        }
        println!("{}", "}");
      },
      Err(Exit::WithMessage(msg)) => println!("Error: {}", msg),
      Err(Exit::Quit) => break,
    }
  }
}


const RADIX: u32 = 10;
const INTEGER_RADIX: Decimal = dec!(10);

fn to_decimal(input: u32) -> Decimal {
  <u32 as Into<Decimal>>::into(input)
}

fn usize_to_decimal(input: usize) -> Decimal {
  <usize as Into<Decimal>>::into(input)
}


type ReadAction = Rc<dyn Fn(char, &mut RPState) -> Result<(), Exit>>;

pub struct ReadTable {
  action_map: BTreeMap<char, ReadAction>,
  default_action: ReadAction,
}

impl ReadTable {
  fn new(default_action: fn(char, &mut RPState) -> Result<(), Exit>)
    -> ReadTable
  {
    ReadTable {
      action_map: BTreeMap::new(),
      default_action: Rc::new(default_action),
    }
  }

  fn set_action(&mut self, c: char,
         action: impl Fn(char, &mut RPState) -> Result<(), Exit> + 'static)
  {
    let _ = self.action_map.insert(c, Rc::new(action));
  }

  fn set_digit_action(&mut self,
                       action: fn(char, &mut RPState) -> Result<(), Exit>)
  {
    for range in ['0' ..= '9', 'a' ..= 'z', 'A' ..= 'Z'] {
      for c in range {
        if c.is_digit(RADIX) {
          let _ = self.action_map.insert(c, Rc::new(action.clone()));
        }
      }
    }
  }
}


struct RPState {
  registers: HashMap<String, Value>,
  stack: Vec<Value>,
  mode: Mode,
  readtables: BTreeMap<Mode, ReadTable>,
  reg_command: String,
  wip_str: String,
  eat_count: u32,
  is_num_negative: bool,
  num: Decimal,
  decimal_offset: Decimal
}


// TODO: Generalized shuffles
// [[a-aa\](shuf)]s[dup]
// [[ab-ba\](shuf)]s[swap]
// [[ab-b\](shuf)]s[nip]
//
// fn shuffle(_state: &mut RPState) -> Result<(), Exit> {
//    Err(err_msg("(shuf) not implemented yet!".into()))
// }

pub fn make_default_readtables() -> BTreeMap<Mode, ReadTable> {
  BTreeMap::from([
    (Mode::Integer, make_integer_readtable()),
    (Mode::Decimal, make_decimal_readtable()),
    (Mode::Str, make_string_readtable()),
    (Mode::CommandChar, make_command_character_readtable()),
    (Mode::CommandNamed, make_command_string_readtable()),
    (Mode::RegisterChar, make_register_character_readtable()),
    (Mode::RegisterStr, make_register_string_readtable()),
  ])
}

fn make_command_string_readtable() -> ReadTable {
  let mut result = ReadTable::new(move |c, state| {
    state.wip_str.push(c);

    Ok(())
  });

  result.set_action(')', move |_, state| {
    match state.wip_str.as_str() {
      // TODO:
      // (times)
      // (while)
      // (?) -- (cond ifTrue ifFalse)
      // [(?)x]s[if-else]
      // [[\](?)x]s[if]
      // (r?) -- show contents of registers
      // Scope out
      word => {
        if state.registers.contains_key(word) {
          match &state.registers[word].clone() {
            Value::Str(eval_body) => {
              state.wip_str.clear();
              state.mode = Mode::CommandChar;
              eval(&eval_body, state)
            }
            _ => {
              state.mode = Mode::CommandChar;
              Err(err_msg(format!(
                  "Unable to execute ({}) as a named word", word)))
            }
          }
        } else {
          state.mode = Mode::CommandChar;
          Err(err_msg(format!(
              "Unable to execute ({}) as a named word", word)))
        }
      }
    }
  });

  result
}


fn make_command_character_readtable() -> ReadTable {
  let mut result = ReadTable::new(move |c, _| {
    panic!("{} isn't implemented yet!", c)
  });

  ////////////////////////////////////////////////////////////
  // Whitespace actions
  //
  result.set_action(' ', move |_, _| Ok(()));

  ////////////////////////////////////////////////////////////
  // Literal actions
  //
  result.set_digit_action(move |c, state| {
    state.mode = Mode::Integer;
    state.num = dec!(0);
    state.num += to_decimal(c.to_digit(10)
        .ok_or_else(|| err_msg(format!("{} isn't a digit", c)))?);
    Ok(())
  });

  result.set_action('_', move |_, state| {
    state.mode = Mode::Integer;
    state.num = dec!(0);
    state.is_num_negative = true;
    Ok(())
  });

  result.set_action('[', move |_, state| {
    state.mode = Mode::Str;
    state.wip_str.clear();
    state.eat_count = 0;
    Ok(())
  });

  result.set_action('(', move |_, state| {
    state.mode = Mode::CommandNamed;
    state.wip_str.clear();
    state.eat_count = 0;
    Ok(())
  });

  ////////////////////////////////////////////////////////////
  // Arithmetic operation actions
  //
  let make_arithmetic_action =
          move |arithmetic: fn(Decimal, Decimal) -> Decimal|
  {
    move |_, state: &mut RPState| {
      if state.stack.len() < 2 {
        return Err(err_msg("Data underflow!".into()));
      }

      let right = state.stack.pop().ok_or_else(|| err_msg(
          "Data underflow!".into()))?;
      let left = state.stack.pop().ok_or_else(|| err_msg(
          "Data underflow!".into()))?;

      match (&left, &right) {
        (Value::Num(left), Value::Num(right)) => {
          let result = arithmetic(*left, *right);

          state.stack.push(Value::Num(result));

          Ok(())
        }
        _ =>  {
          state.stack.push(left);
          state.stack.push(right);

          Err(err_msg("Can't do arithmetic on non-numbers!".into()))
        }
      }
    }
  };

  result.set_action('+', make_arithmetic_action(|left, right| left + right));
  result.set_action('-', make_arithmetic_action(|left, right| left - right));
  result.set_action('/', make_arithmetic_action(|left, right| left / right));
  result.set_action('*', make_arithmetic_action(|left, right| left * right));
  result.set_action('%', make_arithmetic_action(|left, right| left % right));
  result.set_action('^', make_arithmetic_action(|left, right|
      Decimal::pow(left, right)));

  result.set_action('v', move |_, state| {
    match state.stack.pop() {
      Some(Value::Num(d)) => {
        match d.sqrt() {
          Some(n) => {
            state.stack.push(Value::Num(n));
            Ok(())
          }
          None => Err(err_msg("Error attempting square root!".into()))
        }
      }
      Some(v) => Err(err_msg(format!("Invalid attempt to sqrt {}", v))),
      None => Err(err_msg("Data underflow!".into()))
    }
  });

  ////////////////////////////////////////////////////////////
  // Stack operation actions
  //
  result.set_action('d', move |_, state| {
    if state.stack.is_empty() {
      return Err(err_msg("Data underflow!".into()));
    }

    let val = state.stack.pop().unwrap();
    state.stack.push(val.clone());
    state.stack.push(val.clone());

    Ok(())
  });

  result.set_action('c', move |_, state| {
    state.stack.clear();

    Ok(())
  });

  result.set_action(',', move |_, state| {
    if state.stack.is_empty() {
      return Err(err_msg("Data underflow!".into()));
    }

    state.stack.pop();

    Ok(())
  });

  result.set_action('!', move |_, state| {
    match state.stack.pop() {
      Some(Value::Num(d)) => {
        if usize_to_decimal(state.stack.len()) < d {
          Err(err_msg(format!("Stack depth should be at least: {}", d)))
        } else {
          Ok(())
        }
      }
      Some(Value::Str(_)) => {
        Err(err_msg("Cannot assert a string as a stack depth".into()))
      }
      None => {
        Err(err_msg("Data underflow!".into()))
      }
    }
  });

  ////////////////////////////////////////////////////////////
  // Register operation actions
  //
  let register_command_action = move |c: char, state: &mut RPState| {
    state.reg_command = c.into();
    state.mode = Mode::RegisterChar;
    Ok::<(), Exit>(())
  };

  result.set_action('l', register_command_action.clone());
  result.set_action('s', register_command_action.clone());

  ////////////////////////////////////////////////////////////
  // Control flow actions
  //
  result.set_action('x', move |_, state| {
    match state.stack.pop() {
      Some(Value::Str(s)) => eval(&s, state),
      Some(v) => Err(err_msg(format!("Cannot eval {}", v))),
      None => Err(err_msg("Data underflow!".into())),
    }
  });

  ////////////////////////////////////////////////////////////
  // Meta actions
  //
  result.set_action('p', move |_, state| {
    if state.stack.is_empty() {
      return Err(err_msg("Data underflow!".into()));
    }

    print_value(state.stack.last().unwrap());

    Ok(())
  });

  result.set_action('n', move |_, state| {
    if state.stack.is_empty() {
      return Err(err_msg("Data underflow!".into()));
    }

    print_value(&state.stack.pop().unwrap());

    Ok(())
  });

  result.set_action('q', move |_, _| Err(Exit::Quit));

  result.set_action('?', move |_, _| {
    print!("{}", HELP_TEXT);
    Ok(())
  });

  result
}

fn finish_num(c: char, state: &mut RPState) -> Result<(), Exit> {
  // print!("finishing number, negative? {}", state.is_num_negative);
  if state.is_num_negative {
    state.num *= dec!(-1);
  }
  state.stack.push(Value::Num(state.num));
  state.mode = Mode::CommandChar;
  state.is_num_negative = false;
  eval_one(c, state)
}


fn make_integer_readtable() -> ReadTable {
  let mut result = ReadTable::new(finish_num);

  result.set_action('.', move |_, state| {
    state.decimal_offset = dec!(1);
    state.mode = Mode::Decimal;

    Ok(())
  });

  result.set_action('_', move |_, state| {
    state.is_num_negative = true;

    Ok(())
  });

  result.set_digit_action(move |c: char, state: &mut RPState| {
    if c.is_digit(RADIX) {
      state.num *= INTEGER_RADIX;
      state.num += to_decimal(c.to_digit(10).unwrap());
    }

    Ok(())
  });

  result
}


fn make_decimal_readtable() -> ReadTable {
  let mut result = ReadTable::new(finish_num);

  result.set_digit_action(move |c: char, state: &mut RPState| {
    if c.is_digit(RADIX) {
      state.decimal_offset *= dec!(0.1);
      state.num += to_decimal(c.to_digit(10).unwrap())
                   * state.decimal_offset;
    }

    Ok(())
  });

  result
}


fn make_string_readtable() -> ReadTable {
  ReadTable::new(move |c, state| {
    if state.eat_count > 0 {
      state.eat_count -= 1;
      state.wip_str.push(c);
    } else if c == '\\' {
      state.eat_count = 1;
    } else if c != ']' {
      state.wip_str.push(c);
    } else if c == ']' {
      state.mode = Mode::CommandChar;
      state.stack.push(Value::Str(state.wip_str.clone()));
      state.wip_str.clear();
    } else {
      return Err(err_msg("Should Not Get Here!".into()))
    }

    Ok(())
  })
}


fn make_register_string_readtable() -> ReadTable {
  let mut result = ReadTable::new(move |c, state| {
    state.wip_str.push(c);

    Ok(())
  });

  result.set_action(']', move |_, state| {
    state.mode = Mode::CommandChar;

    let name = state.wip_str.clone();
    state.wip_str.clear();

    dispatch_register_command(&name, state)
  });

  result
}


fn make_register_character_readtable() -> ReadTable {
  let mut result = ReadTable::new(move |c, state| {
    state.mode = Mode::CommandChar;

    let name = String::from(c);

    dispatch_register_command(&name, state)
  });

  result.set_action('[', move |_, state| {
    state.mode = Mode::RegisterStr;
    Ok(())
  });

  result
}


fn dispatch_register_command(name: &str, state: &mut RPState)
  -> Result<(), Exit>
{
  match state.reg_command.as_str() {
    "l" => load_from_register(name, state),
    "s" => store_to_register(name, state),

    _ => {
      Err(err_msg(format!(
          "Unsupported register command {}", state.reg_command)))
    }
  }
}


fn store_to_register(name: &str, state: &mut RPState) -> Result<(), Exit> {
  if !state.stack.is_empty() {
    state.registers.insert(name.into(), state.stack.pop().unwrap());

    Ok(())
  } else {
    Err(err_msg(format!(
        "Data underflow attempting to store to register {}", name)))
  }
}


fn load_from_register(name: &str, state: &mut RPState) -> Result<(), Exit> {
  if state.registers.contains_key(name) {
    state.stack.push(state.registers[name].clone());

    Ok(())
  } else {
    Err(err_msg(format!(
        "Register {} is empty and was expected not to be", name)))
  }
}


fn eval(input: &str, state: &mut RPState) -> Result<(), Exit> {
  for (_cpos, c) in input.char_indices() {
    eval_one(c, state)?
  }

  match state.mode {
    Mode::Integer | Mode::Decimal => {
      return finish_num(' ', state)
    },
    _ => { }
  };

  Ok(())
}


fn eval_one(c: char, state: &mut RPState) -> Result<(), Exit> {
  if let Some(readtable) = state.readtables.get_mut(&state.mode) {
    if let Some(action) = readtable.action_map.get_mut(&c) {
      (Rc::clone(&action))(c, state)
    } else {
      (Rc::clone(&readtable.default_action))(c, state)
    }
  } else {
    Err(err_msg("No readtable found for this mode".into()))
  }
}