import { mkdirSync, writeFileSync } from 'node:fs';
import { join } from 'node:path';
import {
  createAgentSession,
  DefaultResourceLoader,
  ModelRuntime,
  SessionManager,
  SettingsManager,
  type AgentSession,
  type RetrySettings,
  type ToolDefinition,
} from '@earendil-works/pi-coding-agent';
import type { PiggyChatContext, PiggyMode } from '@pig/core';
import { assertPigToolBoundary } from '../chat';
import { loadPiggyConfig, type PiggyConfig, type PiggyTurnLimits } from '../config';
import {
  isPiggyModelId,
  piggyDefaultModelId,
  piggyModelCatalogue,
  piggyModelsJsonText,
  PIGGY_PROVIDER_ID,
} from './models';
import { buildPiggySystemPrompt } from './prompt';

export { piggyDefaultModelId, piggyModelCatalogue };

/** A message from an earlier turn, replayed so the conversation continues. */
export interface PiggyHistoryTurn {
  role: 'user' | 'assistant';
  content: string;
}

/** Which ceiling a turn passed, and where it stood when it passed it. */
export interface PiggyTurnBreach {
  limit: 'model_calls' | 'tokens';
  modelCalls: number;
  /** Input plus output over every model call so far. */
  tokens: number;
  /** The ceiling that was passed, in that limit's own units. */
  ceiling: number;
}

/**
 * What a turn has spent, and whether it has spent too much.
 *
 * One of these is created per chat turn and written by two independent
 * counters, on purpose. `installTurnBudget` counts inside the harness loop,
 * which is the only place that can stop the next model call before it is made;
 * the chat server counts the `turn_end` events it already subscribes to, which
 * is the only place that still works if a harness upgrade claims the hook the
 * way it has already claimed `beforeToolCall` and `prepareNextTurnWithContext`.
 * Both report absolute counts to `observeTurn`, so the two readings merge
 * instead of double-counting.
 */
export interface PiggyTurnBudget {
  readonly limits: PiggyTurnLimits;
  modelCalls: number;
  tokens: number;
  /** Set once, by whichever counter saw the ceiling passed first. */
  breach?: PiggyTurnBreach;
  /** A model call was made after the breach: the graceful stop did not hold. */
  overran: boolean;
}

export function createTurnBudget(limits: PiggyTurnLimits): PiggyTurnBudget {
  return { limits, modelCalls: 0, tokens: 0, overran: false };
}

/**
 * Merge one counter's reading of the turn so far.
 *
 * `Math.max` rather than `+=` because the two counters describe the same model
 * calls from two vantage points; adding them would halve the effective ceiling
 * and cut real questions off in the middle.
 */
export function observeTurn(budget: PiggyTurnBudget, modelCalls: number, tokens: number): void {
  const seen = Math.max(budget.modelCalls, modelCalls);
  if (budget.breach) {
    // Another model call after the ceiling was passed. The turn was supposed to
    // have stopped; recording it is how an operator finds out that it did not.
    if (seen > budget.breach.modelCalls) budget.overran = true;
  }
  budget.modelCalls = seen;
  budget.tokens = Math.max(budget.tokens, tokens);
  if (budget.breach) return;
  if (budget.modelCalls >= budget.limits.maxModelCalls) {
    budget.breach = {
      limit: 'model_calls',
      modelCalls: budget.modelCalls,
      tokens: budget.tokens,
      ceiling: budget.limits.maxModelCalls,
    };
    return;
  }
  if (budget.tokens >= budget.limits.maxTurnTokens) {
    budget.breach = {
      limit: 'tokens',
      modelCalls: budget.modelCalls,
      tokens: budget.tokens,
      ceiling: budget.limits.maxTurnTokens,
    };
  }
}

export interface CreatePiggySessionOptions {
  mode: PiggyMode;
  /** Defaults to PIGGY_AGENT_MODEL. Must be in the picker's catalogue. */
  modelId?: string;
  /**
   * Read-only because the chat server holds its tool list as `readonly` and
   * nothing here mutates it; a mutable parameter would force every caller into
   * a defensive copy for no gain.
   */
  tools: readonly ToolDefinition[];
  context?: PiggyChatContext;
  history?: readonly PiggyHistoryTurn[];
  /**
   * The turn's cost ceiling. Optional only so a caller that never prompts — the
   * tool-boundary and prompt tests — need not invent one; every caller that
   * spends money passes it.
   */
  budget?: PiggyTurnBudget;
}

export interface PiggySession {
  session: AgentSession;
  modelId: string;
  systemPrompt: string;
  dispose(): void;
}

/** The messages the agent keeps, as the harness types them. */
type PiggyAgentMessage = AgentSession['agent']['state']['messages'][number];

interface PiggyAgentRuntime {
  modelRuntime: ModelRuntime;
  settingsManager: SettingsManager;
  agentDir: string;
  config: PiggyConfig;
}

/**
 * One runtime per process, behind a promise rather than a value.
 *
 * `ModelRuntime.create` reads files, composes providers and resolves
 * credentials. Doing that per turn would put a filesystem round trip in front
 * of every keystroke in the docked panel; doing it per turn *concurrently* —
 * which is what a plain `if (!runtime)` guard gives you under two simultaneous
 * chats — would build two of them and register the credential twice. Caching
 * the promise makes the second caller await the first construction.
 */
let runtimePromise: Promise<PiggyAgentRuntime> | undefined;

async function piggyAgentRuntime(): Promise<PiggyAgentRuntime> {
  runtimePromise ??= buildAgentRuntime();
  try {
    return await runtimePromise;
  } catch (error) {
    // A failed construction must not be cached: the usual cause is a missing or
    // rejected key, and an operator who fixes the environment and retries
    // should not be served the old failure for the life of the process.
    runtimePromise = undefined;
    throw error;
  }
}

/**
 * What Piggy does when Prime Inference says "please retry shortly".
 *
 * Measured on 2026-08-14, on production, roughly every other turn:
 *
 *     [piggy] chat turn ended in an inference error: 429:
 *     {"message":"Rate limit reached. Please retry shortly.",
 *      "type":"rate_limit_exceeded","code":"rate_limited"}
 *
 * and the reader got `{"type":"error","code":"inference_failed"}` and no answer,
 * while a `curl` a second later succeeded. The endpoint asked us to retry and we
 * did not. `withInferenceRetries` in `apps/piggy/src/provider.ts` still guards
 * the queued worker with exactly this policy — bounded attempts, jittered
 * backoff, `Retry-After` honoured, 429 and 5xx retried and no other 4xx ever —
 * and it was lost for the chat when the harness took over the transport.
 *
 * The seam is the harness's own provider-request retry rather than a loop of
 * ours around `session.prompt()`, and the reason is exactly-once. Read
 * `retryProviderRequest` in `@earendil-works/pi-ai/dist/utils/provider-retry.js`
 * and then its one caller in `dist/api/openai-completions.js:139`: it wraps the
 * creation of the request and nothing else, so every attempt it makes happens
 * BEFORE the first byte of the response has been read. A retry there cannot
 * duplicate a content delta, cannot re-run `pig_log_activity`, and cannot
 * re-apply an approved write, because at that instant none of those has
 * happened. The property is structural rather than policed, which is the only
 * kind worth having when the failure mode is writing a CRM row twice. It also
 * reads `retry-after` and `retry-after-ms`, backs off exponentially with jitter,
 * sleeps on the run's own AbortSignal so a caller hanging up wins immediately,
 * and retries 408, 409, 429 and 5xx and no other status.
 *
 * Measured here, with a stubbed fetch, before any of these values were set:
 * `retryProviderRequest` defaults `maxRetries` to 0 and `getProviderRetrySettings`
 * supplies `undefined`, so the harness made exactly one attempt at every model
 * call. That is the whole bug.
 *
 * `stream` is the second, smaller budget, and it is deliberately not the same
 * number. The harness's session-level auto-retry re-drives a turn that failed
 * AFTER the response started, by discarding the errored assistant message and
 * continuing; that recovers a dropped socket, but it regenerates text the reader
 * has already been shown. Measured, on the same stub: a turn that streamed
 * "Idle is " and then lost the stream came back as "Idle is Idle is $12,000." in
 * the client transcript. So it is kept — a mid-stream drop is the one failure
 * the provider-level retry cannot see — but held to a single attempt, and the
 * chat server refuses the replay outright once anything has been delivered.
 */
export interface PiggyInferenceRetryPolicy {
  /** Attempts at getting a response started, including the first. */
  attempts: number;
  /**
   * Deadline on one attempt.
   *
   * A headers deadline, not a turn deadline: the OpenAI client clears its timer
   * in a `finally` the moment `fetch` resolves (openai@6.26.0 client.js:387-411),
   * so it covers connect and response headers and never the streamed body. That
   * is what makes it safe to set this tight — a legitimately long answer is
   * measured by the stall watchdog's idle clock instead, which restarts on every
   * chunk. 20 seconds is the deadline the hand-rolled chat loop used on the same
   * endpoint for the same reason.
   */
  headersTimeoutMs: number;
  /**
   * The longest `Retry-After` worth honouring.
   *
   * Above this the SDK fails the request immediately and says what was asked
   * for, which is the right answer: three attempts each parked on the SDK's own
   * 60-second default would leave somebody staring at a docked panel for three
   * minutes to be told no. Five seconds twice over is the worst this can add.
   */
  maxRetryDelayMs: number;
  /** Attempts at a turn that failed after the response started, first included. */
  streamAttempts: number;
  /** First backoff for those, doubling per attempt. */
  streamBackoffMs: number;
}

export const PIGGY_INFERENCE_RETRY: PiggyInferenceRetryPolicy = {
  attempts: 4,
  headersTimeoutMs: 20_000,
  maxRetryDelayMs: 5_000,
  streamAttempts: 2,
  streamBackoffMs: 1_500,
};

/**
 * The policy above, in the field names the installed harness actually reads.
 *
 * Exported because it is the only honest way to test this: the values are read
 * by `SettingsManager` and nothing else in PIG, so a test asserts that the
 * installed package hands them back rather than asserting that we wrote an
 * object. That check matters more than it sounds. The obvious place to put a
 * request timeout is the model entry in models.json, and it does nothing there:
 * `ModelDefinitionSchema` in the harness (dist/core/model-config.js:133-147) has
 * no `timeoutMs`, `Model` in `@earendil-works/pi-ai` has no such field, and the
 * only reader is `options.timeoutMs`, which `Agent.createLoopConfig()` never
 * populates. A `timeoutMs` written beside `contextWindow` would validate, load,
 * freeze, and be ignored, with nothing anywhere to say so.
 */
export function piggyAgentSettings(
  policy: PiggyInferenceRetryPolicy = PIGGY_INFERENCE_RETRY,
): NonNullable<Parameters<typeof SettingsManager.inMemory>[0]> {
  const retry: RetrySettings = {
    enabled: policy.streamAttempts > 1,
    maxRetries: Math.max(0, policy.streamAttempts - 1),
    baseDelayMs: policy.streamBackoffMs,
    provider: {
      maxRetries: Math.max(0, policy.attempts - 1),
      maxRetryDelayMs: policy.maxRetryDelayMs,
      timeoutMs: policy.headersTimeoutMs,
    },
  };
  return { retry };
}

async function buildAgentRuntime(): Promise<PiggyAgentRuntime> {
  const config = loadPiggyConfig();
  const agentDir = prepareAgentDir(config.PIGGY_AGENT_DIR);
  const modelsPath = join(agentDir, 'models.json');
  writeFileSync(modelsPath, piggyModelsJsonText(), { mode: 0o600 });

  const modelRuntime = await ModelRuntime.create({
    credentials: new EphemeralCredentialStore(),
    modelsPath,
    // The catalogue is the five models we ship, not whatever the endpoint is
    // advertising this week. A network refresh at startup would make process
    // start depend on api.pinference.ai being reachable, for a list we have
    // already decided.
    allowModelNetwork: false,
  });

  // models.json does NOT resolve environment variable names: writing
  // "apiKey": "PRIME_API_KEY" sends the literal string PRIME_API_KEY as the
  // bearer token and the endpoint answers 401. The credential store is the
  // supported path, and this call is the only one that authenticates Piggy.
  await modelRuntime.setRuntimeApiKey(PIGGY_PROVIDER_ID, config.PRIME_API_KEY);

  return {
    modelRuntime,
    // In-memory settings, because SettingsManager.create writes the chosen
    // model and thinking level back to settings.json. With a model picker per
    // user, that would make one person's choice the process-wide default. It is
    // also the only seam that reaches the harness's HTTP call: the retry budget
    // and the request deadline are read off this object once per model call.
    settingsManager: SettingsManager.inMemory(piggyAgentSettings()),
    agentDir,
    config,
  };
}

function prepareAgentDir(agentDir: string): string {
  // 0o700 because models.json and any session artefact the harness decides to
  // write live here, on a box that also runs the API.
  mkdirSync(agentDir, { recursive: true, mode: 0o700 });
  return agentDir;
}

/**
 * The harness's own credential types, reached through the option that consumes
 * them. `@earendil-works/pi-ai` declares them and is a transitive dependency of
 * the harness rather than one of ours, so importing it by name would be a
 * phantom dependency that breaks the moment the harness re-pins its version.
 */
type PiggyCredentialStore = NonNullable<
  NonNullable<Parameters<typeof ModelRuntime.create>[0]>['credentials']
>;
type PiggyCredential = Awaited<ReturnType<PiggyCredentialStore['read']>>;

/**
 * A credential store that forgets.
 *
 * The key is already in the environment; the default file-backed store would
 * write a second copy of a live Prime platform key into auth.json, which
 * nothing in this repo ever cleans up and nothing rotates. Keeping it in memory
 * means the process holding it is the only thing that has it.
 */
class EphemeralCredentialStore implements PiggyCredentialStore {
  private credential: PiggyCredential;
  private chain: Promise<PiggyCredential> = Promise.resolve(undefined);

  async read(): Promise<PiggyCredential> {
    return this.credential;
  }

  async list(): Promise<readonly { providerId: string; type: 'api_key' }[]> {
    return this.credential ? [{ providerId: PIGGY_PROVIDER_ID, type: 'api_key' }] : [];
  }

  async modify(
    _providerId: string,
    fn: (current: PiggyCredential) => Promise<PiggyCredential>,
  ): Promise<PiggyCredential> {
    // Serialised through a promise chain because the contract requires
    // read-modify-write to be mutually exclusive per provider; two sessions
    // starting at once would otherwise interleave their writes.
    const next = this.chain.then(async () => {
      const updated = await fn(this.credential);
      if (updated !== undefined) this.credential = updated;
      return this.credential;
    });
    this.chain = next.catch(() => undefined);
    return next;
  }

  async delete(): Promise<void> {
    this.credential = undefined;
  }
}

function assertUniqueToolNames(tools: readonly ToolDefinition[]): void {
  const seen = new Set<string>();
  for (const tool of tools) {
    // A duplicate name silently shadows one of the two implementations inside
    // the harness registry, which is how a read tool ends up answering for a
    // write tool of the same name.
    if (seen.has(tool.name)) {
      throw new Error(`Piggy was handed two tools named '${tool.name}'.`);
    }
    seen.add(tool.name);
  }
}

/**
 * The security property of this whole change, checked at runtime.
 *
 * `noTools: 'all'` plus an explicit allowlist should already make this
 * impossible, but "should" is doing a lot of work in a sentence about giving a
 * CRM agent a shell. The harness composes tools from several sources —
 * extensions, skills, built-ins, the allowlist — and a future version that
 * changes the precedence between them would leak silently. Comparing the live
 * tool list to what we handed over turns that into a startup failure.
 */
function assertExactToolSet(session: AgentSession, expected: readonly ToolDefinition[]): void {
  const actual = session.agent.state.tools.map((tool) => tool.name).sort();
  const wanted = expected.map((tool) => tool.name).sort();
  const unexpected = actual.filter((name) => !wanted.includes(name));
  const missing = wanted.filter((name) => !actual.includes(name));
  if (unexpected.length > 0 || missing.length > 0) {
    throw new Error(
      `Piggy's tool set does not match its allowlist. Unexpected: [${unexpected.join(', ')}]. Missing: [${missing.join(', ')}].`,
    );
  }
}

/**
 * The harness's own hook type, reached through the object that owns it, so this
 * file keeps its rule of never importing `@earendil-works/pi-ai` — a transitive
 * dependency — by name.
 */
type ShouldStopAfterTurn = NonNullable<AgentSession['agent']['shouldStopAfterTurn']>;
type ShouldStopContext = Parameters<ShouldStopAfterTurn>[0];

/**
 * The only thing that stops the loop before it buys another model call.
 *
 * `agent-loop.js` is a `while (true)` with four exits: the model stops asking
 * for tools, it errors, the run is aborted, or `shouldStopAfterTurn` returns
 * true. Only the last of those is ours, and it is checked after every turn and
 * before every subsequent request, so returning true here means call N+1 is
 * never made — no tokens, no charge, no latency. Aborting instead would also
 * work, but it would cut the turn off mid-flight and lose the answer the model
 * had already paid for.
 *
 * Counting happens here rather than being read from the chat server because
 * this is the callback the loop makes on the way to spending money: it is
 * handed the assistant message that has just been billed, so nothing can be
 * missed between the provider and the ceiling.
 *
 * Any hook already installed is chained rather than replaced. The harness sets
 * `beforeToolCall` and `prepareNextTurnWithContext` on the same object for its
 * own purposes, and a version that starts using this one would otherwise have
 * its behaviour silently deleted by us.
 */
function installTurnBudget(session: AgentSession, budget: PiggyTurnBudget): void {
  const previous = session.agent.shouldStopAfterTurn;
  let modelCalls = 0;
  let tokens = 0;
  session.agent.shouldStopAfterTurn = async (context, signal) => {
    modelCalls += 1;
    tokens += turnUsage(context);
    observeTurn(budget, modelCalls, tokens);
    if (budget.breach) return true;
    return (await previous?.(context, signal)) === true;
  };
}

/**
 * Input plus output for the model call that has just finished.
 *
 * Input is counted because it is billed and because it is most of the money on
 * a tool-heavy turn: every round trip resends the whole transcript and every
 * tool result so far, so the third call of a turn is several times the size of
 * the first. Shape-checked rather than asserted, for the same reason the chat
 * server checks it: the message union includes types that carry no usage.
 */
function turnUsage(context: ShouldStopContext): number {
  const usage = (context.message as { usage?: { input?: unknown; output?: unknown } }).usage;
  const input = typeof usage?.input === 'number' ? usage.input : 0;
  const output = typeof usage?.output === 'number' ? usage.output : 0;
  return input + output;
}

/**
 * Replays earlier turns into the transcript.
 *
 * The harness starts every in-memory session empty, so without this a second
 * message in the same conversation arrives with no idea what the first one
 * said. Only text is replayed: the tool calls of a previous turn are settled
 * history, and re-presenting them without their results would leave the
 * transcript with dangling calls the provider rejects.
 */
function rehydrateHistory(session: AgentSession, history: readonly PiggyHistoryTurn[]): void {
  if (history.length === 0) return;
  const model = session.agent.state.model;
  const timestamp = Date.now();
  const messages: PiggyAgentMessage[] = history.map((turn) =>
    turn.role === 'user'
      ? { role: 'user', content: turn.content, timestamp }
      : {
          role: 'assistant',
          content: [{ type: 'text', text: turn.content }],
          api: model.api,
          provider: model.provider,
          model: model.id,
          // Zeroed, and deliberately so: this turn was billed when it happened.
          // Carrying its real usage forward would double-count it in the
          // session totals the cost line is drawn from.
          usage: {
            input: 0,
            output: 0,
            cacheRead: 0,
            cacheWrite: 0,
            totalTokens: 0,
            cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
          },
          stopReason: 'stop',
          timestamp,
        },
  );
  session.agent.state.messages = messages;
}

/**
 * Builds a Piggy turn on Prime Agent.
 *
 * Everything the harness would otherwise discover from the filesystem is
 * switched off here, and the loader is reloaded by hand: `createAgentSession`
 * only calls `reload()` on a loader it constructed itself, so a loader passed
 * in that is never reloaded yields the stock coding-assistant prompt with no
 * warning of any kind.
 */
export async function createPiggySession(
  options: CreatePiggySessionOptions,
): Promise<PiggySession> {
  const runtime = await piggyAgentRuntime();
  const modelId = options.modelId ?? runtime.config.PIGGY_AGENT_MODEL;
  if (!isPiggyModelId(modelId)) {
    throw new Error(
      `Model ${modelId} is not in the Piggy catalogue; the picker may only offer ${piggyModelCatalogue()
        .map((option) => option.id)
        .join(', ')}.`,
    );
  }

  const model = runtime.modelRuntime.getModel(PIGGY_PROVIDER_ID, modelId);
  if (!model) {
    throw new Error(
      `Prime Inference did not register model ${modelId}; check apps/piggy/src/agent/models.json.`,
    );
  }

  assertUniqueToolNames(options.tools);
  // The third gate, behind `noTools: 'all'` and the explicit allowlist. It is
  // the only one written in PIG's own code, so it is the only one a harness
  // upgrade cannot quietly change the meaning of.
  assertPigToolBoundary(options.tools);

  const systemPrompt = buildPiggySystemPrompt({
    mode: options.mode,
    context: options.context,
    tools: options.tools,
  });
  const loader = new DefaultResourceLoader({
    cwd: runtime.agentDir,
    agentDir: runtime.agentDir,
    settingsManager: runtime.settingsManager,
    noExtensions: true,
    noSkills: true,
    noPromptTemplates: true,
    noThemes: true,
    noContextFiles: true,
    // systemPromptOverride takes the literal text; the `systemPrompt` option is
    // a file source, and handing it a prompt loads nothing and says nothing.
    systemPromptOverride: () => systemPrompt,
    appendSystemPromptOverride: () => [],
  });
  await loader.reload();

  const toolNames = options.tools.map((tool) => tool.name);
  const { session } = await createAgentSession({
    agentDir: runtime.agentDir,
    cwd: runtime.agentDir,
    modelRuntime: runtime.modelRuntime,
    // The per-turn budget is applied to the model rather than the request
    // because the harness reads the ceiling off the model it is given. Clamped
    // to the model's own maximum so raising the budget cannot ask for more
    // than the endpoint will return.
    model: { ...model, maxTokens: Math.min(runtime.config.PIGGY_AGENT_MAX_TOKENS, model.maxTokens) },
    settingsManager: runtime.settingsManager,
    thinkingLevel: runtime.config.PIGGY_AGENT_THINKING,
    noTools: 'all',
    tools: toolNames,
    customTools: [...options.tools],
    sessionManager: SessionManager.inMemory(),
    resourceLoader: loader,
  });

  assertExactToolSet(session, options.tools);
  if (options.budget) installTurnBudget(session, options.budget);
  rehydrateHistory(session, options.history ?? []);

  let disposed = false;
  return {
    session,
    modelId,
    systemPrompt,
    dispose: () => {
      if (disposed) return;
      disposed = true;
      // Abort before dispose: a session disposed mid-turn keeps the upstream
      // inference socket open and billing, because dropping the listeners does
      // not tell the provider to stop generating.
      void session.abort().catch(() => {});
      session.dispose();
    },
  };
}
