(* Evaluator: matching, substitution, rewriting driver *)

open Ast

module Env = Map.Make (String)

type env = node Env.t

let is_register = function
  | Atom s when String.length s > 0 && s.[0] = '?' -> true
  | _ -> false

let reg_name = function
  | Atom s -> String.sub s 1 (String.length s - 1)
  | _ -> invalid_arg "reg_name"

let rec equal_node a b =
  match (a, b) with
  | Atom x, Atom y -> String.equal x y
  | List xs, List ys ->
      List.length xs = List.length ys
      && List.for_all2 equal_node xs ys
  | _ -> false

(* Attempt to match pattern against tree, accumulating bindings in env *)
let rec match_pattern env pattern tree =
  match pattern with
  | p when is_register p -> (
      let name = reg_name p in
      match Env.find_opt name env with
      | None -> Ok (Env.add name tree env)
      | Some bound -> if equal_node bound tree then Ok env else Error ())
  | Atom x -> (
      match tree with Atom y when String.equal x y -> Ok env | _ -> Error ())
  | List pxs -> (
      match tree with
      | List txs when List.length pxs = List.length txs ->
          List.fold_left2
            (fun acc p t -> match acc with Ok e -> match_pattern e p t | _ -> acc)
            (Ok env) pxs txs
      | _ -> Error ())

let rec substitute env = function
  | a when is_register a -> (
      let name = reg_name a in
      match Env.find_opt name env with
      | Some v -> v
      | None -> a)
  | Atom _ as a -> a
  | List xs -> List (List.map (substitute env) xs)

(* Configuration for evaluation behavior *)
type config = { mutable quiet : bool; mutable allow_access : bool; mutable cycles : int }

let default_config () = { quiet = false; allow_access = false; cycles = 0x200000 }

(* Devices and dynamic forms (AST-based) *)

let rec node_to_string = function
  | Atom s -> s
  | List xs -> String.concat "" (List.map node_to_string xs)

let parse_int_atom s =
  try
    if String.length s > 1 && s.[0] = '#' then int_of_string_opt ("0x" ^ String.sub s 1 (String.length s - 1))
    else int_of_string_opt s
  with _ -> None

let rec flatten_atoms acc = function
  | Atom s -> s :: acc
  | List xs -> List.fold_left flatten_atoms acc xs

let device_eval ~(cfg : config) (node : node) : node option =
  match node with
  | List (Atom "?:" :: args) ->
      (* Print or ALU: if first arg is an operator [+ - * / % & ^ | = ! > <], perform arithmetic; else print all args *)
      let alu_ops = [ "+"; "-"; "*"; "/"; "%"; "&"; "^"; "|"; "="; "!"; ">"; "<" ] in
      let result_opt =
        match args with
        | Atom op :: rest when List.mem op alu_ops ->
            let nums =
              List.filter_map
                (fun a ->
                  match a with Atom s -> parse_int_atom s | List _ -> parse_int_atom (node_to_string a))
                rest
            in
            (match nums with
            | [] -> None
            | acc :: tl ->
                let f acc n =
                  match op with
                  | "+" -> acc + n
                  | "-" -> acc - n
                  | "*" -> acc * n
                  | "/" -> (try acc / n with _ -> acc)
                  | "%" -> (try acc mod n with _ -> acc)
                  | "&" -> acc land n
                  | "^" -> acc lxor n
                  | "|" -> acc lor n
                  | "=" -> if acc = n then 1 else 0
                  | "!" -> if acc <> n then 1 else 0
                  | ">" -> if acc > n then 1 else 0
                  | "<" -> if acc < n then 1 else 0
                  | _ -> acc
                in
                let res = List.fold_left f acc tl in
                Some (Atom (string_of_int res)))
        | _ ->
            let s = String.concat "" (List.map node_to_string args) in
            if not cfg.quiet then output_string stdout s;
            Some (List [])
      in
      result_opt
  | List [ Atom "?_"; path ] ->
      if not cfg.allow_access then Some (Atom "NAF")
      else (
        let filepath = node_to_string path in
        try
          let ch = open_in filepath in
          let len = in_channel_length ch in
          let buf = really_input_string ch len in
          close_in ch;
          (* Parse the imported string as a node; on failure, return Atom *)
          try Some (Parse.parse buf) with _ -> Some (Atom buf)
        with _ -> Some (Atom "NAF"))
  | List [ Atom "?~" ] ->
      if not cfg.allow_access then Some (Atom "EOF")
      else (
        try
          let buf = really_input_string stdin (in_channel_length stdin) in
          Some (Atom buf)
        with _ -> Some (Atom "EOF"))
  | List [ Atom "?^"; arg ] ->
      let atoms = List.rev (flatten_atoms [] arg) in
      Some (Atom (String.concat "" atoms))
  | List [ Atom "?."; arg ] -> (
      match arg with List xs -> Some (List xs) | _ -> Some arg)
  | List [ Atom "?*"; arg ] -> (
      match arg with
      | Atom s -> Some (List (List.init (String.length s) (fun i -> Atom (String.make 1 s.[i]))))
      | List xs -> Some (List xs))
  | _ -> None

let is_define = function List [ Atom "<>"; _; _ ] -> true | _ -> false
let is_undefine = function List [ Atom "><"; _; _ ] -> true | _ -> false

let extract_lr = function List [ _; l; r ] -> (l, r) | _ -> failwith "extract_lr"

(* Walk positions in a tree in preorder, returning a zipper-like path *)
type path = (node list * node list) list

let nodes_with_paths tree =
  let rec go acc path = function
    | Atom _ as a -> (List.rev path, a) :: acc
    | List xs as l ->
        let acc' = (List.rev path, l) :: acc in
        let rec children acc left right =
          match right with
          | [] -> acc
          | n :: rs ->
              let path' = (left, rs) :: path in
              let acc'' = go acc path' n in
              children acc'' (n :: left) rs
        in
        children acc' [] xs
  in
  List.rev (go [] [] tree)

let replace_at_path root path new_subtree =
  let rec rebuild = function
    | [] -> new_subtree
    | (left, right) :: rest ->
        let rebuilt_child = rebuild rest in
        let lst = List.rev left @ (rebuilt_child :: right) in
        List lst
  in
  match path with [] -> new_subtree | _ -> rebuild (List.rev path)

let try_dynamic ~(cfg : config) ~(rules : rule list ref) tree =
  let candidates = nodes_with_paths tree in
  let rec loop = function
    | [] -> None
    | (path, node) :: rest -> (
        (* Devices *)
        match device_eval ~cfg node with
        | Some replacement -> Some (replace_at_path tree path replacement)
        | None ->
            (* Define / Undefine *)
            if is_define node then (
              let l, r = extract_lr node in
              rules := !rules @ [ { left = l; right = r } ];
              Some (replace_at_path tree path (List []))
            ) else if is_undefine node then (
              let l, r = extract_lr node in
              rules := List.filter (fun ru -> not (equal_node ru.left l && equal_node ru.right r)) !rules;
              Some (replace_at_path tree path (List []))
            ) else loop rest)
  in
  loop candidates

let try_apply_rule rule tree =
  let candidates = nodes_with_paths tree in
  let rec loop = function
    | [] -> None
    | (path, node) :: rest -> (
        match match_pattern Env.empty rule.left node with
        | Ok env ->
            let rhs = substitute env rule.right in
            Some (replace_at_path tree path rhs)
        | Error () -> loop rest)
  in
  loop candidates

let eval ?(config = default_config ()) rules tree =
  let rules_ref = ref rules in
  let rec loop cycles_left current =
    if cycles_left <= 0 then current
    else
      match try_dynamic ~cfg:config ~rules:rules_ref current with
      | Some t' -> loop (cycles_left - 1) t'
      | None ->
          let rec scan = function
            | [] -> None
            | r :: rs -> (
                match try_apply_rule r current with Some t' -> Some t' | None -> scan rs)
          in
          (match scan !rules_ref with Some t' -> loop (cycles_left - 1) t' | None -> current)
  in
  loop config.cycles tree