package com.gzzm.lobster.runtime;

import com.gzzm.lobster.llm.CancelReason;
import com.gzzm.lobster.llm.RunCancelledException;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * RunContext —— 一次 run 的执行上下文 / Execution context of a single run.
 *
 * <p>把过去散落在 {@link AgentRuntime} 主循环里的几个局部变量
 *（{@code turn} / {@code cancelled} / {@code pendingRequestId} / 用量统计）
 * 归集成一个对象，便于：
 * <ul>
 *   <li>工具层按 {@link #cancelFlag} 传播取消</li>
 *   <li>{@link RunEventBus} 按 runId 分发事件</li>
 *   <li>未来把主循环搬到 executor 时不必重新发明上下文容器</li>
 * </ul>
 *
 * <p>设计原则（与 Claude Code 对齐）：<b>不做事件溯源</b>，RunContext 只是运行时状态容器；
 * 持久化事实一律在 transcript / run 表。所以这里不保存事件历史，只保存"当前"状态。
 *
 * <p>Java 8：无 record —— 用 public final + getter 混合风格。
 */
public final class RunContext {

    public final String runId;
    public final String threadId;
    public final String userId;
    public final String orgId;
    public final int    maxTurns;
    public volatile long deadlineAtMs;     // 0 表示无 deadline
    public final AtomicReference<CancelReason> cancel = new AtomicReference<CancelReason>();
    public final AtomicInteger turn = new AtomicInteger(0);
    public final UsageCounter usage = new UsageCounter();
    public final RunEventBus events;       // 进程内，不落库

    public RunContext(String runId, String threadId, String userId, String orgId,
                      int maxTurns, long deadlineAtMs, RunEventBus events) {
        this.runId = runId;
        this.threadId = threadId;
        this.userId = userId;
        this.orgId = orgId;
        this.maxTurns = maxTurns;
        this.deadlineAtMs = deadlineAtMs;
        this.events = events == null ? InMemoryRunEventBus.NOOP : events;
    }

    /** 请求硬取消 / Request hard cancel with a reason. */
    public void requestCancel(CancelReason reason) {
        if (reason == null) reason = CancelReason.USER;
        cancel.compareAndSet(null, reason);
    }

    public boolean isCancelled() { return cancel.get() != null; }
    public CancelReason cancelReason() { return cancel.get(); }

    public void resetDeadline(long deadlineAtMs) {
        this.deadlineAtMs = deadlineAtMs;
    }

    /**
     * 有边界检查的"请继续"探测 / Throw {@link RunCancelledException} if cancelled or
     * deadline passed.
     *
     * <p>主循环在每个可取消边界调一次：进入/退出 turn、工具调用前后、LLM 流式 poll。
     */
    public void checkCancelled() {
        CancelReason r = cancel.get();
        if (r != null) throw new RunCancelledException(r);
        if (deadlineAtMs > 0 && System.currentTimeMillis() > deadlineAtMs) {
            cancel.compareAndSet(null, CancelReason.TIMEOUT);
            throw new RunCancelledException(CancelReason.TIMEOUT, "run deadline exceeded");
        }
    }

    /** token / 工具调用计数器 —— 便于 BUDGET 类取消和审计。 */
    public static final class UsageCounter {
        public final AtomicInteger promptTokens = new AtomicInteger();
        public final AtomicInteger completionTokens = new AtomicInteger();
        public final AtomicInteger toolCalls = new AtomicInteger();

        public void addPromptTokens(int n) { if (n > 0) promptTokens.addAndGet(n); }
        public void addCompletionTokens(int n) { if (n > 0) completionTokens.addAndGet(n); }
        public void incToolCalls() { toolCalls.incrementAndGet(); }
    }
}