## bridge_loop.R — R-side REPL for the Python r_bridge module.
## Reads sentinel-prefixed JSON requests from stdin, writes responses to stdout.

# ---------------------------------------------------------------------------
# 1. Bootstrap
# ---------------------------------------------------------------------------

options(warn = 1, OutDec = ".")
# Redirect message() output to stderr so it never touches the protocol stdout
options(useFancyQuotes = FALSE)
# On Windows, ensure stdout/stderr use UTF-8 to match the Python side.
if (.Platform$OS.type == "windows") {
  tryCatch(
    {
      # R >= 4.2 UCRT: sets stdout/stderr to UTF-8 binary mode (no-op on UCRT, needed on MSVCRT)
      invisible(Sys.setlocale("LC_CTYPE", "en_US.UTF-8"))
    },
    error = function(e) invisible(NULL)
  )
}

if (!requireNamespace("jsonlite", quietly = TRUE)) {
  message("r_bridge: jsonlite not found, attempting to install...")
  tryCatch(
    install.packages("jsonlite", repos = "https://cloud.r-project.org", quiet = TRUE),
    error = function(e) message("r_bridge: install.packages failed: ", conditionMessage(e))
  )
}

if (!requireNamespace("jsonlite", quietly = TRUE)) {
  cat('__RBRIDGE__:{"id":"__init__","status":"error","payload":{"message":"jsonlite is not installed and could not be installed automatically. Run install.packages(\\"jsonlite\\") in R."},"warnings":[]}\n')
  flush(stdout())
  quit(save = "no", status = 2)
}

SENTINEL <- "__RBRIDGE__:"

send_response <- function(id, status, payload, warnings = list()) {
  obj <- list(id = id, status = status, payload = payload, warnings = as.list(warnings))
  line <- jsonlite::toJSON(obj, auto_unbox = TRUE, null = "null", na = "null", digits = 15)
  cat(SENTINEL, line, "\n", sep = "")
  flush(stdout())
}

# Signal ready
send_response("__init__", "ok", list(ready = TRUE, jsonlite_version = as.character(packageVersion("jsonlite")), arrow_available = FALSE))

# ---------------------------------------------------------------------------
# 2. R → JSON serializer helpers
# ---------------------------------------------------------------------------

r_to_payload <- function(x) {
  if (is.null(x)) return(NULL)

  if (is.data.frame(x))   return(encode_dataframe(x))
  if (is.factor(x))       return(encode_factor(x))
  if (is.matrix(x))       return(encode_matrix(x))
  if (inherits(x, "POSIXct") || inherits(x, "POSIXlt")) return(encode_posixct(x))
  if (inherits(x, "Date")) return(list(`__type__` = "date", iso8601 = format(x, "%Y-%m-%d")))

  if (is.list(x)) {
    if (!is.null(names(x))) {
      return(list(`__type__` = "named_list", value = lapply(x, r_to_payload)))
    } else {
      return(list(`__type__` = "list", value = lapply(x, r_to_payload)))
    }
  }

  if (is.character(x)) {
    if (length(x) == 1) {
      if (is.na(x)) return(list(`__type__` = "na", na_type = "character"))
      return(x)
    }
    vals <- lapply(x, function(v) if (is.na(v)) list(`__type__` = "na", na_type = "character") else v)
    return(list(`__type__` = "character_vector", value = vals))
  }

  if (is.logical(x)) {
    if (length(x) == 1) {
      if (is.na(x)) return(list(`__type__` = "na", na_type = "logical"))
      return(x)
    }
    vals <- lapply(x, function(v) if (is.na(v)) list(`__type__` = "na", na_type = "logical") else v)
    return(list(`__type__` = "logical_vector", value = vals))
  }

  if (is.integer(x)) {
    if (length(x) == 1) {
      if (is.na(x)) return(list(`__type__` = "na", na_type = "integer"))
      return(list(`__type__` = "integer", value = x))
    }
    vals <- lapply(x, function(v) if (is.na(v)) list(`__type__` = "na", na_type = "integer") else list(`__type__` = "integer", value = v))
    return(list(`__type__` = "integer_vector", value = vals))
  }

  if (is.complex(x)) {
    if (length(x) == 1) return(list(`__type__` = "complex", r = Re(x), i = Im(x)))
    vals <- lapply(x, function(v) list(`__type__` = "complex", r = Re(v), i = Im(v)))
    return(list(`__type__` = "complex_vector", value = vals))
  }

  if (is.numeric(x)) {
    if (length(x) == 1) return(encode_scalar_numeric(x))
    vals <- lapply(x, encode_scalar_numeric)
    return(list(`__type__` = "numeric_vector", value = vals))
  }

  # Fallback: coerce to character
  as.character(x)
}

encode_scalar_numeric <- function(v) {
  if (is.na(v))       return(list(`__type__` = "na", na_type = "real"))
  if (is.nan(v))      return(list(`__type__` = "nan"))
  if (is.infinite(v)) return(list(`__type__` = "inf", sign = if (v > 0) 1 else -1))
  # Tag explicitly so Python can reliably decode as float, not int
  list(`__type__` = "numeric_scalar", value = v)
}

encode_factor <- function(x) {
  list(
    `__type__` = "factor",
    levels     = as.list(levels(x)),
    codes      = as.list(as.integer(x)),   # 1-based, NA → NA_integer_
    ordered    = is.ordered(x)
  )
}

encode_dataframe <- function(x) {
  cols <- lapply(x, r_to_payload)
  list(`__type__` = "dataframe", columns = cols)
}

encode_matrix <- function(x) {
  list(
    `__type__` = "matrix",
    nrow  = nrow(x),
    ncol  = ncol(x),
    data  = as.list(as.vector(t(x))),   # row-major (byrow=TRUE on Python side)
    byrow = TRUE
  )
}

encode_posixct <- function(x) {
  list(`__type__` = "posixct", iso8601 = format(as.POSIXct(x, tz = "UTC"), "%Y-%m-%dT%H:%M:%S"))
}

# ---------------------------------------------------------------------------
# 3. JSON → R deserializer helpers
# ---------------------------------------------------------------------------

payload_to_r <- function(x) {
  if (is.null(x)) return(NULL)
  if (!is.list(x) || is.null(x[["__type__"]])) {
    # Plain scalar or already-R value
    return(x)
  }

  tag <- x[["__type__"]]

  switch(tag,
    "integer"          = x$value,
    "nan"              = NaN,
    "inf"              = if (x$sign > 0) Inf else -Inf,
    "complex"          = complex(real = x$r, imaginary = x$i),
    "posixct"          = as.POSIXct(x$iso8601, tz = "UTC"),
    "date"             = as.Date(x$iso8601),
    "na"               = decode_na(x),
    "numeric_vector"   = decode_numeric_vector(x),
    "integer_vector"   = decode_integer_vector(x),
    "logical_vector"   = decode_logical_vector(x),
    "character_vector" = decode_character_vector(x),
    "complex_vector"   = decode_complex_vector(x),
    "matrix"           = decode_matrix(x),
    "array"            = decode_array(x),
    "dataframe"        = decode_dataframe(x),
    "factor"           = decode_factor(x),
    "named_list"       = decode_named_list(x),
    "list"             = lapply(x$value, payload_to_r),
    stop(paste("r_bridge: unknown type tag:", tag))
  )
}

decode_na <- function(x) {
  switch(x$na_type,
    "integer"   = NA_integer_,
    "real"      = NA_real_,
    "complex"   = NA_complex_,
    "character" = NA_character_,
    "logical"   = NA
  )
}

decode_numeric_vector <- function(x) {
  sapply(x$value, function(v) {
    if (is.list(v) && !is.null(v[["__type__"]])) payload_to_r(v) else as.numeric(v)
  })
}

decode_integer_vector <- function(x) {
  sapply(x$value, function(v) {
    if (is.list(v) && !is.null(v[["__type__"]])) payload_to_r(v) else as.integer(v)
  })
}

decode_logical_vector <- function(x) {
  sapply(x$value, function(v) {
    if (is.list(v) && !is.null(v[["__type__"]])) payload_to_r(v) else as.logical(v)
  })
}

decode_character_vector <- function(x) {
  sapply(x$value, function(v) {
    if (is.list(v) && !is.null(v[["__type__"]])) payload_to_r(v) else as.character(v)
  })
}

decode_complex_vector <- function(x) {
  sapply(x$value, function(v) {
    if (is.list(v) && !is.null(v[["__type__"]])) payload_to_r(v) else as.complex(v)
  })
}

decode_matrix <- function(x) {
  data <- sapply(x$data, function(v) if (is.list(v)) payload_to_r(v) else v)
  if (isTRUE(x$byrow)) {
    matrix(data, nrow = x$nrow, ncol = x$ncol, byrow = TRUE)
  } else {
    matrix(data, nrow = x$nrow, ncol = x$ncol)
  }
}

decode_array <- function(x) {
  data  <- sapply(x$data, function(v) if (is.list(v)) payload_to_r(v) else v)
  array(data, dim = unlist(x$shape))
}

decode_dataframe <- function(x) {
  cols <- lapply(x$columns, payload_to_r)
  as.data.frame(cols, stringsAsFactors = FALSE)
}

decode_factor <- function(x) {
  levels <- unlist(x$levels)
  codes  <- unlist(x$codes)
  vals   <- ifelse(codes >= 1, levels[codes], NA_character_)
  factor(vals, levels = levels, ordered = isTRUE(x$ordered))
}

decode_named_list <- function(x) {
  lapply(x$value, payload_to_r)
}

# ---------------------------------------------------------------------------
# 4. Operation handlers
# ---------------------------------------------------------------------------

handle_set_var <- function(payload) {
  name  <- payload$name
  value <- payload_to_r(payload$value)
  assign(name, value, envir = .GlobalEnv)
  list(status = "ok", payload = list())
}

handle_call_func <- function(payload) {
  func_name <- payload$func
  args      <- lapply(payload$args %||% list(), payload_to_r)
  kwargs    <- lapply(payload$kwargs %||% list(), payload_to_r)
  fn        <- tryCatch(match.fun(func_name), error = function(e) stop(paste("Function not found:", func_name)))
  result    <- do.call(fn, c(args, kwargs))
  list(status = "ok", payload = r_to_payload(result))
}

handle_eval_expr <- function(payload) {
  result <- eval(parse(text = payload$expr), envir = .GlobalEnv)
  list(status = "ok", payload = r_to_payload(result))
}

handle_get_var <- function(payload) {
  name  <- payload$name
  value <- tryCatch(get(name, envir = .GlobalEnv), error = function(e) stop(paste("Variable not found:", name)))
  list(status = "ok", payload = r_to_payload(value))
}

handle_ls <- function(payload) {
  list(status = "ok", payload = list(names = as.list(ls(envir = .GlobalEnv))))
}

handle_exec_script <- function(payload) {
  code <- payload[["code"]]
  con  <- textConnection(code)
  on.exit(close(con))
  source(con, local = FALSE, echo = FALSE)
  list(status = "ok", payload = NULL)
}

# Null-coalescing operator
`%||%` <- function(a, b) if (!is.null(a)) a else b

# ---------------------------------------------------------------------------
# 5. Main REPL loop
# ---------------------------------------------------------------------------

stdin_con <- file("stdin", open = "r")

repeat {
  line <- readLines(con = stdin_con, n = 1, warn = FALSE)
  if (length(line) == 0) break   # EOF — Python process died

  if (!startsWith(line, SENTINEL)) next   # stray output, ignore

  body <- substring(line, nchar(SENTINEL) + 1)
  req  <- tryCatch(
    jsonlite::fromJSON(body, simplifyVector = FALSE),
    error = function(e) NULL
  )

  if (is.null(req)) {
    send_response("__unknown__", "error",
      list(message = "Failed to parse request JSON"), list())
    next
  }

  id <- req$id %||% "__unknown__"
  op <- req$op %||% ""

  warnings_caught <- character(0)

  result <- withCallingHandlers(
    tryCatch(
      switch(op,
        ping     = list(status = "ok", payload = list(pong = TRUE)),
        shutdown = {
          send_response(id, "ok", list())
          quit(save = "no", status = 0)
        },
        set_var   = handle_set_var(req$payload),
        call_func = handle_call_func(req$payload),
        eval_expr = handle_eval_expr(req$payload),
        get_var     = handle_get_var(req$payload),
        ls          = handle_ls(req$payload),
        exec_script = handle_exec_script(req$payload),
        list(status = "error", payload = list(message = paste("Unknown op:", op)))
      ),
      error = function(e) {
        tb <- tryCatch(paste(capture.output(traceback()), collapse = "\n"), error = function(e2) "")
        list(status = "error", payload = list(
          message   = conditionMessage(e),
          call      = paste(deparse(conditionCall(e)), collapse = ""),
          traceback = tb
        ))
      }
    ),
    warning = function(w) {
      warnings_caught <<- c(warnings_caught, conditionMessage(w))
      invokeRestart("muffleWarning")
    }
  )

  send_response(id, result$status, result$payload, warnings_caught)
}
