# Context Engine
# R as the thalamic filter: decides what reaches the LLM's attention.
#
# Target: pack ~100K tokens (400KB) of the best possible context every turn.
#
# Components:
#   .context_engine$conversation  - data.frame: queryable conversation index
#   .context_engine$file_index    - named list: path -> lines (in-memory project)
#   .context_engine$symbol_index  - saber symbols() output, cached
#   .context_engine$payload       - pre-computed context, ready to splice
#   .context_engine$cwd           - project working directory
#   .context_engine$config        - config snapshot

.context_engine <- new.env(parent = emptyenv())

# Lifecycle ----

#' Initialize the context engine
#'
#' Loads file index, saber symbols, and prepares for payload assembly.
#'
#' @param cwd Project working directory
#' @param config Config list from load_config()
#' @return Invisible NULL
#' @noRd
ce_init <- function(cwd, config = list()) {
    .context_engine$cwd <- cwd
    .context_engine$config <- config
    .context_engine$payload <- NULL
    .context_engine$dirty <- TRUE

    # Initialize conversation index
    .context_engine$conversation <- data.frame(turn = integer(),
        role = character(), content = character(), tokens = integer(),
        tool_calls = I(list()), files_touched = I(list()),
        timestamp = as.POSIXct(character()), stringsAsFactors = FALSE)

    # Build in-memory file index
    ce_index_files(cwd)

    # Cache saber symbols (best-effort)
    ce_update_symbols(cwd)

    invisible(NULL)
}

#' Shut down the context engine
#'
#' @return Invisible NULL
#' @noRd
ce_shutdown <- function() {
    invisible(NULL)
}

# Conversation indexing ----

#' Index a conversation turn
#'
#' Adds a turn to the in-memory conversation data.frame.
#'
#' @param turn Integer turn number
#' @param role "user" or "assistant"
#' @param content Message text
#' @param tool_calls Character vector of tool names called (NULL for none)
#' @param files_touched Character vector of file paths touched (NULL for none)
#' @return Invisible turn number
#' @noRd
ce_index_turn <- function(turn, role, content, tool_calls = NULL,
                          files_touched = NULL) {
    tokens <- ceiling(nchar(content) / 4L)

    row <- data.frame(turn = as.integer(turn), role = role,
                      content = content, tokens = tokens,
                      tool_calls = I(list(tool_calls %||% character())),
                      files_touched = I(list(files_touched %||% character())),
                      timestamp = Sys.time(), stringsAsFactors = FALSE)

    conv <- .context_engine$conversation
    .context_engine$conversation <- rbind(conv, row)
    .context_engine$dirty <- TRUE

    invisible(turn)
}

#' Get the conversation data.frame
#'
#' @return data.frame with turn, role, content, tokens, tool_calls,
#'   files_touched, timestamp
#' @noRd
ce_conversation <- function() {
    .context_engine$conversation
}

#' Search conversation by keyword
#'
#' @param query Character string to search for
#' @return data.frame subset of matching turns
#' @noRd
ce_search_conversation <- function(query) {
    conv <- .context_engine$conversation
    if (nrow(conv) == 0) {
        return(conv)
    }

    hits <- grep(query, conv$content, ignore.case = TRUE)
    conv[hits,, drop = FALSE]
}

#' Total tokens in conversation
#'
#' @return Integer
#' @noRd
ce_conversation_tokens <- function() {
    conv <- .context_engine$conversation
    if (nrow(conv) == 0) {
        return(0L)
    }
    sum(conv$tokens)
}

# File index ----

#' Index project files into memory
#'
#' Reads all text files matching patterns into a named list for instant grep.
#'
#' @param cwd Project directory
#' @param patterns Glob patterns to match
#' @param max_file_size Max file size in bytes (skip larger files)
#' @return Invisible count of indexed files
#' @noRd
ce_index_files <- function(cwd,
                           extensions = c("R", "r", "Rmd", "md", "json", "yaml", "yml", "c",
        "cpp", "h", "js", "css", "html", "sql", "sh", "py",
        "Rd"),
                           extra_files = c("DESCRIPTION", "NAMESPACE", "Makefile", "Dockerfile", ".Rbuildignore"),
                           max_file_size = 100000L) {
    index <- list()
    # winslash="/" everywhere keeps index keys free of backslashes, so
    # regex prefix stripping and cross-platform glob/regex matching work
    # identically on Windows and POSIX.
    cwd_norm <- normalizePath(cwd, winslash = "/", mustWork = FALSE)

    # Recursive list of all files, then filter by extension
    ext_pattern <- sprintf("\\.(%s)$", paste(extensions, collapse = "|"))
    all_files <- list.files(cwd, recursive = TRUE, full.names = TRUE)

    if (endsWith(cwd_norm, "/")) {
        prefix <- cwd_norm
    } else {
        prefix <- paste0(cwd_norm, "/")
    }
    for (f in all_files) {
        norm <- normalizePath(f, winslash = "/", mustWork = FALSE)
        rel <- if (startsWith(norm, prefix)) {
            substring(norm, nchar(prefix) + 1L)
        } else if (identical(norm, cwd_norm)) {
            ""
        } else {
            norm
        }

        # Skip .git, node_modules, renv
        if (grepl("^(\\.git|node_modules|renv)/", rel)) {
            next
        }

        # Accept by extension or exact name
        base <- basename(f)
        if (!grepl(ext_pattern, base, ignore.case = TRUE) &&
            !base %in% extra_files) {
            next
        }

        # Skip binary/oversized
        sz <- tryCatch(file.size(f), error = function(e) Inf)
        if (sz > max_file_size) {
            next
        }

        lines <- tryCatch(readLines(f, warn = FALSE), error = function(e) NULL)
        if (!is.null(lines)) {
            index[[rel]] <- lines
        }
    }

    .context_engine$file_index <- index
    invisible(length(index))
}

#' Update specific files in the index
#'
#' Re-reads only the specified files. Call after tool writes.
#'
#' @param paths Character vector of file paths (absolute or relative)
#' @return Invisible count of updated files
#' @noRd
ce_update_files <- function(paths) {
    cwd <- .context_engine$cwd %||% getwd()
    updated <- 0L

    for (p in paths) {
        # An absolute path starts with "/" on POSIX or a drive letter
        # followed by colon on Windows (e.g. "C:/..." or "C:\...").
        is_abs <- startsWith(p, "/") ||
        grepl("^[A-Za-z]:[/\\\\]", p)
        if (is_abs) {
            abs_path <- p
        } else {
            abs_path <- file.path(cwd, p)
        }
        if (!file.exists(abs_path)) {
            # File deleted: remove from index
            # Can't normalizePath a deleted file, but we can normalize its
            # parent dir (which still exists) and append the basename. This
            # handles macOS where /var -> /private/var: the stored key used
            # fully-normalized paths, so the deletion key must match.
            parent <- dirname(abs_path)
            norm_parent <- normalizePath(parent, winslash = "/",
                mustWork = FALSE)
            norm_abs <- paste(norm_parent, basename(abs_path), sep = "/")
            norm_cwd <- normalizePath(cwd, winslash = "/", mustWork = FALSE)
            if (endsWith(norm_cwd, "/")) {
                prefix <- norm_cwd
            } else {
                prefix <- paste0(norm_cwd, "/")
            }
            rel <- if (startsWith(norm_abs, prefix)) {
                substring(norm_abs, nchar(prefix) + 1L)
            } else {
                norm_abs
            }
            if (nchar(rel) > 0) {
                .context_engine$file_index[[rel]] <- NULL
            }
            next
        }

        lines <- tryCatch(readLines(abs_path, warn = FALSE),
                          error = function(e) NULL)
        if (!is.null(lines)) {
            norm_cwd <- normalizePath(cwd, winslash = "/", mustWork = FALSE)
            norm_abs <- normalizePath(abs_path, winslash = "/", mustWork = FALSE)
            if (endsWith(norm_cwd, "/")) {
                prefix <- norm_cwd
            } else {
                prefix <- paste0(norm_cwd, "/")
            }
            rel <- if (startsWith(norm_abs, prefix)) {
                substring(norm_abs, nchar(prefix) + 1L)
            } else {
                norm_abs
            }
            .context_engine$file_index[[rel]] <- lines
            updated <- updated + 1L
        }
    }

    .context_engine$dirty <- TRUE
    invisible(updated)
}

#' In-memory grep across file index
#'
#' Search all indexed files for a pattern. Microseconds, no disk IO.
#'
#' @param pattern Regex pattern
#' @param file_glob Optional glob to filter files (e.g., "*.R")
#' @return data.frame with file, line_number, text columns
#' @noRd
ce_grep <- function(pattern, file_glob = NULL) {
    index <- .context_engine$file_index
    if (is.null(index) || length(index) == 0) {
        return(data.frame(file = character(), line_number = integer(),
                          text = character(), stringsAsFactors = FALSE))
    }

    # Filter files by glob if specified
    files <- names(index)
    if (!is.null(file_glob)) {
        # Convert glob to regex
        glob_regex <- utils::glob2rx(file_glob)
        files <- files[grepl(glob_regex, basename(files))]
    }

    results <- list()
    for (f in files) {
        lines <- index[[f]]
        hits <- grep(pattern, lines)
        if (length(hits) > 0) {
            results <- c(results, list(data.frame(
                        file = f,
                        line_number = hits,
                        text = lines[hits],
                        stringsAsFactors = FALSE
                    )))
        }
    }

    if (length(results) == 0) {
        return(data.frame(file = character(), line_number = integer(),
                          text = character(), stringsAsFactors = FALSE))
    }
    do.call(rbind, results)
}

#' Get file contents from index
#'
#' @param path Relative file path
#' @return Character vector of lines, or NULL if not indexed
#' @noRd
ce_file <- function(path) {
    if (is.null(path) || length(path) == 0 || nchar(path) == 0) {
        return(NULL)
    }
    idx <- .context_engine$file_index
    if (is.null(idx) || !path %in% names(idx)) {
        return(NULL)
    }
    idx[[path]]
}

#' Count of indexed files and total lines
#'
#' @return Named integer vector: files, lines, chars
#' @noRd
ce_file_stats <- function() {
    index <- .context_engine$file_index
    if (is.null(index) || length(index) == 0) {
        return(c(files = 0L, lines = 0L, chars = 0L))
    }
    n_lines <- sum(vapply(index, length, integer(1)))
    n_chars <- sum(vapply(index, function(x) sum(nchar(x)), numeric(1)))
    c(files = length(index), lines = n_lines, chars = as.integer(n_chars))
}

# Symbol/AST layer (via saber) ----

#' Update symbol index from saber
#'
#' @param cwd Project directory
#' @return Invisible logical (TRUE if updated)
#' @noRd
ce_update_symbols <- function(cwd) {
    if (!requireNamespace("saber", quietly = TRUE)) {
        .context_engine$symbol_index <- NULL
        return(invisible(FALSE))
    }

    .context_engine$symbol_index <- tryCatch(saber::symbols(cwd),
        error = function(e) NULL)
    invisible(!is.null(.context_engine$symbol_index))
}

#' Get function definitions from symbol index
#'
#' @return data.frame with name, file, line, exported columns, or NULL
#' @noRd
ce_definitions <- function() {
    idx <- .context_engine$symbol_index
    if (is.null(idx)) {
        return(NULL)
    }
    idx$defs
}

#' Get function call graph from symbol index
#'
#' @return data.frame with caller, callee, file, line columns, or NULL
#' @noRd
ce_call_graph <- function() {
    idx <- .context_engine$symbol_index
    if (is.null(idx)) {
        return(NULL)
    }
    idx$calls
}

#' Find callers of a function across the project
#'
#' @param fn Function name
#' @return data.frame of callers, or empty data.frame
#' @noRd
ce_blast_radius <- function(fn) {
    if (!requireNamespace("saber", quietly = TRUE)) {
        return(data.frame(caller = character(), project = character(),
                          file = character(), line = integer(),
                          stringsAsFactors = FALSE))
    }
    cwd <- .context_engine$cwd %||% getwd()
    tryCatch(
             saber::blast_radius(fn, project = cwd),
             error = function(e) {
        data.frame(caller = character(), project = character(),
                   file = character(), line = integer(),
                   stringsAsFactors = FALSE)
    }
    )
}

#' Find relevant code for a prompt using symbol index
#'
#' Matches prompt keywords against function names in the symbol index,
#' then returns the files and surrounding code.
#'
#' @param prompt User prompt
#' @param max_results Max number of matching functions to return
#' @return Character string of relevant code snippets
#' @noRd
ce_related_code <- function(prompt, max_results = 10L) {
    defs <- ce_definitions()
    if (is.null(defs) || nrow(defs) == 0) {
        return("")
    }

    # Extract keywords from prompt
    words <- unique(tolower(strsplit(prompt, "[^a-zA-Z0-9_]+")[[1]]))
    words <- words[nchar(words) > 2]
    if (length(words) == 0) {
        return("")
    }

    # Score each definition by keyword overlap
    scores <- vapply(seq_len(nrow(defs)), function(i) {
        name <- tolower(defs$name[i])
        # Split function name on _ and . for partial matching
        name_parts <- strsplit(name, "[_.]")[[1]]
        sum(vapply(words, function(w) {
            if (grepl(w, name, fixed = TRUE)) return(1)
            if (any(grepl(w, name_parts, fixed = TRUE))) return(0.5)
            0
        }, numeric(1)))
    }, numeric(1))

    # Get top matches
    top_idx <- head(order(scores, decreasing = TRUE), max_results)
    top_idx <- top_idx[scores[top_idx] > 0]
    if (length(top_idx) == 0) {
        return("")
    }

    # Retrieve code snippets from file index
    snippets <- character()
    for (i in top_idx) {
        file <- defs$file[i]
        line <- defs$line[i]
        fn_name <- defs$name[i]

        lines <- ce_file(file)
        if (is.null(lines)) {
            next
        }

        # Get function body: start at definition, read until next definition
        # or end of reasonable block
        start <- max(1, line)
        end <- min(length(lines), line + 30)

        # Try to find end of function (simple heuristic: closing brace at col 1)
        for (j in (line + 1):min(length(lines), line + 100)) {
            if (grepl("^\\}", lines[j])) {
                end <- j
                break
            }
        }

        snippet <- paste(sprintf("%s:%d-%d %s()", file, start, end, fn_name),
                         paste(lines[start:end], collapse = "\n"), sep = "\n")
        snippets <- c(snippets, snippet)
    }

    paste(snippets, collapse = "\n\n")
}

# Payload assembly ----

#' Compute the full context payload
#'
#' Assembles system prompt + project context + workspace state,
#' targeting ~100K tokens total.
#'
#' @param prompt Current user prompt (for relevance scoring)
#' @param system_base Base system prompt (from load_context)
#' @param tools_json Tool definitions as JSON string
#' @return List with system (character), tokens_used (integer)
#' @noRd
ce_compute_payload <- function(prompt, system_base, tools_json = "") {
    target_tokens <- 100000L

    # Fixed costs
    system_base_tokens <- ceiling(nchar(system_base) / 4L)
    tool_tokens <- ceiling(nchar(tools_json) / 4L)
    conv_tokens <- ce_conversation_tokens()

    # Budget remaining for dynamic context
    budget_tokens <- target_tokens - system_base_tokens - tool_tokens -
    conv_tokens
    budget_tokens <- max(budget_tokens, 5000L) # floor: always have some budget
    budget_chars <- budget_tokens * 4L

    # Assemble dynamic context, ranked by relevance
    parts <- character()

    # 1. File tree (always: tiny, high value, prevents tool-based exploration)
    tree <- ce_file_tree()
    if (nchar(tree) > 0) {
        parts <- c(parts, "## Project Files", "", "```", tree, "```", "")
    }

    # 2. saber briefing (if available, most bang per token)
    briefing <- ce_get_briefing()
    if (nchar(briefing) > 0) {
        parts <- c(parts, "## Project Briefing", "", briefing, "")
    }

    # 3. Key file contents (high-priority files the LLM should know upfront)
    key_budget <- min(budget_chars %/% 3, 60000L)
    key_context <- ce_key_files(prompt, max_chars = key_budget)
    if (nchar(key_context) > 0) {
        parts <- c(parts, "## Key Files", "", key_context, "")
    }

    # 4. Relevant code from symbol index
    related <- ce_related_code(prompt, max_results = 8L)
    if (nchar(related) > 0) {
        parts <- c(parts, "## Relevant Code", "", related, "")
    }

    # 5. Files recently touched in conversation
    recent_files <- ce_recent_files(n = 5L)
    if (length(recent_files) > 0) {
        file_context <- ce_format_files(recent_files,
                                        max_chars = budget_chars %/% 3)
        if (nchar(file_context) > 0) {
            parts <- c(parts, "## Recently Touched Files", "", file_context, "")
        }
    }

    # 6. Workspace state
    ws_budget <- min(budget_chars %/% 4, 40000L)
    ws_context <- ws_format_context(
                                    ws_retrieve(prompt, budget_chars = ws_budget)
    )
    if (nchar(ws_context) > 0) {
        parts <- c(parts, ws_context, "")
    }

    # Combine and enforce budget
    dynamic_context <- paste(parts, collapse = "\n")
    if (nchar(dynamic_context) > budget_chars) {
        dynamic_context <- substr(dynamic_context, 1, budget_chars)
    }

    enriched_system <- if (nchar(dynamic_context) > 0) {
        paste(system_base, "\n\n", dynamic_context)
    } else {
        system_base
    }

    tokens_used <- ceiling(nchar(enriched_system) / 4L) + tool_tokens +
    conv_tokens

    list(
         system = enriched_system,
         tokens_used = tokens_used,
         budget_remaining = target_tokens - tokens_used
    )
}

#' Get saber briefing for current project
#'
#' @return Character string (empty if saber not available)
#' @noRd
ce_get_briefing <- function() {
    if (!requireNamespace("saber", quietly = TRUE)) {
        return("")
    }

    cwd <- .context_engine$cwd %||% getwd()
    project <- basename(cwd)

    tryCatch({
        text <- saber::briefing(project = project)
        if (is.character(text) && nchar(text) > 0) text else ""
    }, error = function(e) "")
}

#' Get files recently touched in conversation
#'
#' @param n Max number of files to return
#' @return Character vector of relative file paths
#' @noRd
ce_recent_files <- function(n = 5L) {
    conv <- .context_engine$conversation
    if (nrow(conv) == 0) {
        return(character())
    }

    # Collect all files_touched from recent turns, most recent first
    all_files <- character()
    for (i in rev(seq_len(nrow(conv)))) {
        ft <- conv$files_touched[[i]]
        if (length(ft) > 0) {
            all_files <- c(all_files, ft)
        }
    }

    # Unique, preserving order (most recent first)
    unique_files <- unique(all_files)
    head(unique_files, n)
}

#' Format file contents for context injection
#'
#' @param paths Character vector of relative file paths
#' @param max_chars Max total characters
#' @return Character string
#' @noRd
ce_format_files <- function(paths, max_chars = 20000L) {
    parts <- character()
    used <- 0L

    for (p in paths) {
        lines <- ce_file(p)
        if (is.null(lines)) {
            next
        }

        content <- paste(lines, collapse = "\n")
        if (used + nchar(content) > max_chars) {
            # Truncate this file to fit
            remaining <- max_chars - used
            if (remaining < 200) {
                break
            }
            content <- paste0(substr(content, 1, remaining - 20),
                              "\n... (truncated)")
        }

        parts <- c(parts, sprintf("### %s\n```\n%s\n```", p, content))
        used <- used + nchar(content)
    }

    paste(parts, collapse = "\n\n")
}

#' Generate a compact file tree from the index
#'
#' @return Character string, one path per line
#' @noRd
ce_file_tree <- function() {
    idx <- .context_engine$file_index
    if (is.null(idx) || length(idx) == 0) {
        return("")
    }
    paths <- sort(names(idx))
    # Add line counts for each file
    tree_lines <- vapply(paths, function(p) {
        n <- length(idx[[p]])
        sprintf("%s (%d lines)", p, n)
    }, character(1))
    paste(tree_lines, collapse = "\n")
}

#' Inject key project files into context
#'
#' Prioritizes files that give the LLM immediate understanding of the project:
#' DESCRIPTION, README, entry points (app.R, server.R, main.R), config files.
#' Remaining budget goes to small R source files ranked by size (small first).
#'
#' @param prompt User prompt (for relevance scoring)
#' @param max_chars Max total characters
#' @return Character string of file contents
#' @noRd
ce_key_files <- function(prompt, max_chars = 60000L) {
    idx <- .context_engine$file_index
    if (is.null(idx) || length(idx) == 0) {
        return("")
    }

    paths <- names(idx)

    # Priority tiers (order matters within tiers)
    tier1 <- c("DESCRIPTION", "README.md", "README.Rmd")
    tier2 <- c("app.R", "server.R", "ui.R", "main.R", "global.R",
               "NAMESPACE", "AGENTS.md", "PLAN.md")

    # Score each file: tier membership + prompt keyword match + size penalty
    prompt_words <- unique(tolower(strsplit(prompt, "[^a-zA-Z0-9_]+")[[1]]))
    prompt_words <- prompt_words[nchar(prompt_words) > 2]

    scores <- vapply(paths, function(p) {
        base <- basename(p)
        score <- 0

        # Tier bonuses
        if (base %in% tier1) score <- score + 100
        if (base %in% tier2) score <- score + 50
        if (grepl("^R/", p)) score <- score + 10

        # Prompt keyword match against filename
        p_lower <- tolower(p)
        for (w in prompt_words) {
            if (grepl(w, p_lower, fixed = TRUE)) {
                score <- score + 20
            }
        }

        # Prefer smaller files (more files fit in budget)
        n_chars <- sum(nchar(idx[[p]])) + length(idx[[p]])
        score <- score - (n_chars / 1000)

        score
    }, numeric(1))

    # Sort by score descending
    ordered <- paths[order(scores, decreasing = TRUE)]

    # Pack files into budget
    parts <- character()
    used <- 0L

    for (p in ordered) {
        lines <- idx[[p]]
        content <- paste(lines, collapse = "\n")
        nc <- nchar(content)

        if (used + nc > max_chars) {
            # Try truncating if file is high priority
            remaining <- max_chars - used
            if (remaining > 500 && basename(p) %in% c(tier1, tier2)) {
                content <- paste0(substr(content, 1, remaining - 20),
                                  "\n... (truncated)")
                nc <- nchar(content)
            } else if (remaining < 200) {
                break
            } else {
                next
            }
        }

        parts <- c(parts, sprintf("### %s\n```\n%s\n```", p, content))
        used <- used + nc
    }

    paste(parts, collapse = "\n\n")
}

#' Get pre-computed payload (fast path)
#'
#' Returns cached payload if available, otherwise NULL.
#'
#' @return List with system and tokens_used, or NULL
#' @noRd
ce_get_payload <- function() {
    .context_engine$payload
}

#' Fast re-rank with actual prompt
#'
#' If pre-computed payload exists, does a fast adjustment.
#' Otherwise, computes from scratch.
#'
#' @param prompt User prompt
#' @param system_base Base system prompt
#' @param tools_json Tool definitions as JSON
#' @return List with system and tokens_used
#' @noRd
ce_rerank <- function(prompt, system_base, tools_json = "") {
    # For now, always compute fresh. Pre-computed payload is a future
    # optimization once we verify the synchronous path is fast enough.
    ce_compute_payload(prompt, system_base, tools_json)
}

# Utility ----

#' Extract tool names from an llm.api result
#'
#' @param result Result from llm.api::agent()
#' @return Character vector of tool names called
#' @noRd
ce_extract_tool_calls <- function(result) {
    if (is.null(result) || is.null(result$history)) {
        return(character())
    }

    tools <- character()
    for (msg in result$history) {
        if (!is.null(msg$tool_calls)) {
            for (tc in msg$tool_calls) {
                tools <- c(tools, tc$name %||% tc[["function"]]$name)
            }
        }
    }
    unique(tools)
}

#' Extract file paths touched by tool calls
#'
#' Heuristic: look for file path arguments in tool calls.
#'
#' @param result Result from llm.api::agent()
#' @return Character vector of file paths
#' @noRd
ce_extract_files_touched <- function(result) {
    if (is.null(result) || is.null(result$history)) {
        return(character())
    }

    files <- character()
    for (msg in result$history) {
        if (!is.null(msg$tool_calls)) {
            for (tc in msg$tool_calls) {
                args <- tc$arguments %||% tc[["function"]]$arguments
                if (is.character(args)) {
                    args <- tryCatch(
                                     jsonlite::fromJSON(args, simplifyVector = FALSE),
                                     error = function(e) list()
                    )
                }
                # Common file path argument names
                for (key in c("path", "con", "file", "filename")) {
                    if (!is.null(args[[key]]) && is.character(args[[key]])) {
                        files <- c(files, args[[key]])
                    }
                }
            }
        }
    }
    unique(files)
}

