import { readFileSync } from "node:fs";
import { getPackageDir, VERSION } from "@earendil-works/pi-coding-agent";
import { homedir } from "node:os";
import { dirname, join, resolve } from "node:path";

const EXTENSION_ID = "pi-system-prompt-patcher";
const SETTINGS_FILE = `${EXTENSION_ID}.json`;

export type JsonPrimitive = boolean | null | number | string;

export type JsonValue = JsonObject | JsonPrimitive | JsonValue[];

export type JsonObject = {
  [key: string]: JsonValue;
};

export type PromptPatcherContext = {
  model: { provider: string; id: string } | undefined;
  hasUI: boolean;
  abort: () => void;
  ui: {
    notify: (message: string, level: "error") => void;
  };
};

export type BeforeProviderRequestHandler = (
  event: { payload: unknown },
  ctx: PromptPatcherContext,
) => JsonObject | undefined;

export type PromptPatcherApi = {
  on: (event: "before_provider_request", handler: BeforeProviderRequestHandler) => void;
};

type Replacement = {
  target: string;
  replacement: string;
};

type ExpandedReplacement = Replacement & {
  configuredTarget: string;
};

/** Runtime values that rule targets and replacements can reference by placeholder. */
export type RuntimeValues = {
  piPackageDir: string;
  piVersion: string;
};

/**
 * Replace `{piPackageDir}` with the running Pi's package directory and `{piVersion}` with its
 * version, so one rule matches every Pi installation instead of a single install path.
 */
export function expandPlaceholders(text: string, values: RuntimeValues): string {
  return text
    .replaceAll("{piPackageDir}", values.piPackageDir)
    .replaceAll("{piVersion}", values.piVersion);
}

type ProviderSettings = {
  replacementFile?: string;
  models: Record<string, string>;
};

type Settings = {
  providers: Record<string, ProviderSettings>;
};

export const MINIMUM_PI_VERSION = "1.0.1";

/** Compare dotted numeric release versions; prerelease suffixes rank below their release. */
export function atLeastVersion(version: unknown, minimum: string): boolean {
  if (typeof version !== "string") return false;
  const match = /^(\d+)\.(\d+)\.(\d+)(-.+)?$/.exec(version);
  if (!match) return false;
  const actual = [Number(match[1]), Number(match[2]), Number(match[3])];
  const required = minimum.split(".").map(Number);
  for (let index = 0; index < 3; index++) {
    const left = actual[index] ?? 0;
    const right = required[index] ?? 0;
    if (left !== right) return left > right;
  }
  return match[4] === undefined;
}

/**
 * Refuse to load on a Pi runtime older than the release this version targets.
 * Pi installs packages without resolving peer dependencies, so the peer range
 * alone does not stop an older Pi from loading the extension.
 */
export function requirePiVersion(extension: string, version: unknown = VERSION): void {
  if (atLeastVersion(version, MINIMUM_PI_VERSION)) return;
  const reported = typeof version === "string" ? `Pi ${version}` : "no Pi version";
  throw new Error(
    `${extension} requires Pi ${MINIMUM_PI_VERSION} or later, but the running Pi reports ` +
      `${reported}. Exit Pi and start Pi ${MINIMUM_PI_VERSION} or later. ` +
      "/reload cannot upgrade the running runtime.",
  );
}

export default function systemPromptPatcher(pi: PromptPatcherApi) {
  requirePiVersion(EXTENSION_ID);
  pi.on("before_provider_request", (event, ctx) => {
    if (!ctx.model || !hasSystemInstructions(event.payload)) {
      return;
    }

    const settingsPath = resolveSettingsPath();
    const settings = loadSettings(settingsPath, ctx);
    if (!settings) {
      return;
    }

    const configuredPath = selectReplacementFile(settings, ctx.model.provider, ctx.model.id);
    if (!configuredPath) {
      return;
    }

    const replacementPath = resolveConfiguredPath(configuredPath, settingsPath);
    const replacements = loadReplacements(replacementPath, ctx);
    if (!replacements) {
      return;
    }

    // Pi resolves its documentation paths the same way when it builds the system prompt.
    const values = { piPackageDir: resolve(getPackageDir()), piVersion: VERSION };
    const expanded = replacements.map(({ target, replacement }) => ({
      target: expandPlaceholders(target, values),
      replacement: expandPlaceholders(replacement, values),
      configuredTarget: target,
    }));
    return patchSystemPrompt(event.payload, expanded, replacementPath, ctx);
  });
}

function patchSystemPrompt(
  payload: JsonObject,
  replacements: ExpandedReplacement[],
  replacementPath: string,
  ctx: PromptPatcherContext,
): JsonObject | undefined {
  const system = payload["system"];

  if (Object.hasOwn(payload, "system") && !isString(system) && !Array.isArray(system)) {
    reportError(
      ctx,
      `${EXTENSION_ID}: provider payload did not contain a supported system prompt.`,
    );
    return undefined;
  }

  let patched = payload;

  for (const [index, { target, replacement, configuredTarget }] of replacements.entries()) {
    let found = false;
    patched = mapSystemPromptText(patched, (text) => {
      if (text.includes(target)) found = true;
      return text.replaceAll(target, replacement);
    });

    if (!found) {
      reportMissingTarget(ctx, index, target, configuredTarget, replacementPath);
      return undefined;
    }
  }

  return patched;
}

function mapSystemPromptText(payload: JsonObject, transform: (text: string) => string): JsonObject {
  const patched = { ...payload };
  const system = payload["system"];
  if (system !== undefined) patched["system"] = mapPromptContent(system, transform);

  const messages = payload["messages"];
  if (Array.isArray(messages)) {
    patched["messages"] = messages.map((message) => {
      if (!isJsonObject(message) || message["role"] !== "system") return message;
      const content = message["content"];
      if (content === undefined) return message;
      return { ...message, content: mapPromptContent(content, transform) };
    });
  }
  return patched;
}

function mapPromptContent(content: JsonValue, transform: (text: string) => string): JsonValue {
  if (isString(content)) return transform(content);
  if (!Array.isArray(content)) return content;
  return content.map((block) => {
    if (!isJsonObject(block) || block["type"] !== "text" || !isString(block["text"])) return block;
    return { ...block, text: transform(block["text"]) };
  });
}

function reportMissingTarget(
  ctx: PromptPatcherContext,
  index: number,
  target: string,
  configuredTarget: string,
  replacementPath: string,
) {
  const lines = [
    `${EXTENSION_ID}: replacement ${index + 1} target was not found in the request system prompt.`,
    "No prompt replacements were applied for this request.",
    "Aborting the current agent turn.",
    `Replacement file: ${replacementPath}`,
    `Missing target: ${JSON.stringify(target)}`,
  ];
  if (configuredTarget !== target) {
    lines.push(`Configured target: ${JSON.stringify(configuredTarget)}`);
  }
  reportError(ctx, lines.join("\n"));
  ctx.abort();
}

function loadSettings(settingsPath: string, ctx: PromptPatcherContext): Settings | undefined {
  let parsed: unknown;

  try {
    parsed = JSON.parse(readFileSync(settingsPath, "utf8"));
  } catch (error) {
    if (!(error instanceof Error)) {
      throw error;
    }
    reportError(ctx, `${EXTENSION_ID}: failed to read ${settingsPath}: ${formatError(error)}`);
    return;
  }

  if (!isJsonObject(parsed) || !isJsonObject(parsed["providers"])) {
    reportError(ctx, `${EXTENSION_ID}: ${settingsPath} must contain a providers object.`);
    return;
  }

  const providers: Record<string, ProviderSettings> = {};

  for (const [provider, value] of Object.entries(parsed["providers"])) {
    if (provider.length === 0) {
      reportError(ctx, `${EXTENSION_ID}: provider names in ${settingsPath} must not be empty.`);
      return;
    }

    if (!isJsonObject(value)) {
      reportError(ctx, `${EXTENSION_ID}: provider ${JSON.stringify(provider)} must be an object.`);
      return;
    }

    let replacementFile: string | undefined;
    if (Object.hasOwn(value, "replacementFile")) {
      if (!isString(value["replacementFile"]) || value["replacementFile"].length === 0) {
        reportError(
          ctx,
          `${EXTENSION_ID}: provider ${JSON.stringify(provider)} must have a non-empty string replacementFile.`,
        );
        return;
      }
      replacementFile = value["replacementFile"];
    }

    const models: Record<string, string> = {};
    if (Object.hasOwn(value, "models")) {
      if (!isJsonObject(value["models"])) {
        reportError(
          ctx,
          `${EXTENSION_ID}: provider ${JSON.stringify(provider)} models must be an object.`,
        );
        return;
      }

      for (const [model, modelReplacementFile] of Object.entries(value["models"])) {
        if (model.length === 0) {
          reportError(
            ctx,
            `${EXTENSION_ID}: model names for provider ${JSON.stringify(provider)} must not be empty.`,
          );
          return;
        }
        if (!isString(modelReplacementFile) || modelReplacementFile.length === 0) {
          reportError(
            ctx,
            `${EXTENSION_ID}: model ${JSON.stringify(model)} for provider ${JSON.stringify(provider)} must map to a non-empty replacement file path.`,
          );
          return;
        }
        models[model] = modelReplacementFile;
      }
    }

    if (replacementFile === undefined && Object.keys(models).length === 0) {
      reportError(
        ctx,
        `${EXTENSION_ID}: provider ${JSON.stringify(provider)} must configure replacementFile or at least one model.`,
      );
      return;
    }

    providers[provider] = replacementFile === undefined ? { models } : { replacementFile, models };
  }

  return { providers };
}

function selectReplacementFile(
  settings: Settings,
  provider: string,
  model: string,
): string | undefined {
  const providerSettings = settings.providers[provider];
  if (!providerSettings) {
    return;
  }

  return providerSettings.models[model] ?? providerSettings.replacementFile;
}

function resolveConfiguredPath(configuredPath: string, settingsPath: string): string {
  if (configuredPath === "~") {
    return homedir();
  }
  if (/^~[\\/]/.test(configuredPath)) {
    return resolve(homedir(), configuredPath.slice(2));
  }
  return resolve(dirname(settingsPath), configuredPath);
}

function loadReplacements(
  replacementPath: string,
  ctx: PromptPatcherContext,
): Replacement[] | undefined {
  let parsed: unknown;

  try {
    parsed = JSON.parse(readFileSync(replacementPath, "utf8"));
  } catch (error) {
    if (!(error instanceof Error)) {
      throw error;
    }
    reportError(ctx, `${EXTENSION_ID}: failed to read ${replacementPath}: ${formatError(error)}`);
    return;
  }

  if (!Array.isArray(parsed)) {
    reportError(ctx, `${EXTENSION_ID}: ${replacementPath} must contain an array.`);
    return;
  }

  const replacementItems: unknown[] = parsed;
  const replacements: Replacement[] = [];

  for (const [index, item] of replacementItems.entries()) {
    if (!isJsonObject(item)) {
      reportError(ctx, `${EXTENSION_ID}: replacement ${index + 1} must be an object.`);
      return;
    }

    const { target, replacement } = item;
    if (!isString(target) || target.length === 0) {
      reportError(
        ctx,
        `${EXTENSION_ID}: replacement ${index + 1} must have a non-empty string target.`,
      );
      return;
    }

    if (!isString(replacement)) {
      reportError(ctx, `${EXTENSION_ID}: replacement ${index + 1} must have a string replacement.`);
      return;
    }

    replacements.push({ target, replacement });
  }

  return replacements;
}

function resolveSettingsPath(): string {
  const agentDir = process.env["PI_CODING_AGENT_DIR"] ?? join(homedir(), ".pi", "agent");
  return join(agentDir, SETTINGS_FILE);
}

function reportError(ctx: PromptPatcherContext, message: string) {
  console.error(message);
  if (ctx.hasUI) {
    ctx.ui.notify(message, "error");
  }
}

function formatError(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

function isString(value: unknown): value is string {
  return typeof value === "string";
}

function isJsonValue(value: unknown): value is JsonValue {
  if (
    value === null ||
    typeof value === "boolean" ||
    typeof value === "number" ||
    typeof value === "string"
  ) {
    return true;
  }
  if (Array.isArray(value)) return value.every(isJsonValue);
  return typeof value === "object" && Object.values(value).every(isJsonValue);
}

function isJsonObject(value: unknown): value is JsonObject {
  return typeof value === "object" && isJsonValue(value) && value !== null && !Array.isArray(value);
}

function hasSystemInstructions(value: unknown): value is JsonObject {
  if (!isJsonObject(value)) return false;
  if (Object.hasOwn(value, "system")) return true;
  const messages = value["messages"];
  return (
    Array.isArray(messages) &&
    messages.some((message) => isJsonObject(message) && message["role"] === "system")
  );
}