/**
 * pi-params — per-provider per-model parameter overrides for pi.dev
 *
 * Reads config from .pi/params.toml and injects parameter overrides
 * (temperature, top_p, top_k, frequency_penalty, presence_penalty,
 *  repetition_penalty, max_tokens, stop) into the provider request
 * payload before it's sent to the model.
 *
 * Also provides runtime slash commands:
 *   /params              — show current overrides
 *   /params set <k> <v>  — set a session override
 *   /params unset <k>    — remove a session override
 *   /params preset <n>   — apply a named preset
 *   /params reset        — clear all session overrides
 *
 * Works with any provider (OpenAI, Anthropic, Google, OpenRouter, etc.)
 * by injecting params into the provider-agnostic request body.
 *
 * Config format: .pi/params.toml
 *   ["<provider>"."<model>"]
 *   temperature = 0.7
 *   top_p = 0.9
 *
 * Wildcards: "*" matches any provider or model.
 * Priority: session override > exact match > provider-star > star-model > "*"
 */

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { readFileSync } from "node:fs";
import { join } from "node:path";
import { parse } from "smol-toml";

// ── Types ────────────────────────────────────────────────────

/** Raw parsed config shape: { provider: { model: params } } */
type RawConfig = Record<string, Record<string, Record<string, unknown>>>;

/** Tracks the currently active provider and model */
interface CurrentModel {
  provider: string | null;
  model: string | null;
}

/** Known parameter keys we support injecting */
const SUPPORTED_KEYS = new Set([
  "temperature",
  "top_p",
  "top_k",
  "frequency_penalty",
  "presence_penalty",
  "repetition_penalty",
  "max_tokens",
  "stop",
]);

// ── Presets ──────────────────────────────────────────────────

const PRESETS: Record<string, Record<string, unknown>> = {
  creative: { temperature: 1.0, top_p: 0.95 },
  precise: { temperature: 0.1, top_p: 0.1 },
  balanced: { temperature: 0.5, top_p: 0.8 },
  code: { temperature: 0.2, top_p: 0.5 },
  brainstorm: { temperature: 1.2, top_p: 0.95, frequency_penalty: 0.5 },
  long: { max_tokens: 16384 },
};

// ── Module-level state ───────────────────────────────────────

let config: RawConfig | null = null;
let current: CurrentModel = { provider: null, model: null };

/** Runtime overrides set via /params commands (highest priority) */
let sessionOverrides: Record<string, unknown> = {};

// ── Config loader ────────────────────────────────────────────

/**
 * Load and parse the params.toml config file.
 * Returns null (with a console.warn) if the file is missing or malformed.
 */
function loadConfig(cwd: string): RawConfig | null {
  const filePath = join(cwd, ".pi", "params.toml");

  try {
    const raw = readFileSync(filePath, "utf-8");
    const parsed = parse(raw) as unknown as RawConfig;

    // Basic shape validation: must be an object with provider → model → params
    if (typeof parsed !== "object" || parsed === null) {
      console.warn("[pi-params] config is not an object, ignoring");
      return null;
    }

    return parsed;
  } catch (err: unknown) {
    if (err && typeof err === "object" && "code" in err && (err as { code: string }).code === "ENOENT") {
      console.warn("[pi-params] no .pi/params.toml found — running without overrides");
    } else {
      const msg = err instanceof Error ? err.message : String(err);
      console.warn(`[pi-params] failed to parse .pi/params.toml: ${msg}`);
    }
    return null;
  }
}

// ── Wildcard matching ────────────────────────────────────────

/**
 * Match specificity levels (higher = more specific):
 *   3 = exact provider + exact model
 *   2 = exact provider + wildcard model (provider-star)
 *   1 = wildcard provider + exact model (star-model)
 *   0 = wildcard + wildcard ("*")
 */
type MatchResult = {
  params: Record<string, unknown>;
  specificity: number; // 0-3
};

/**
 * Find matching config entries for a given provider + model and merge them
 * in specificity order (most specific overrides least specific).
 * Returns null if no match found at all.
 *
 * Merge order (higher overrides lower):
 *   1. Global defaults (*)
 *   2. Model wildcard across providers (star/model)
 *   3. Provider wildcard (provider-star)
 *   4. Exact provider + exact model (provider/model)
 *
 * Session overrides (from /params commands) are applied on top
 * of this result in the request handler.
 */
function findBestMatch(
  cfg: RawConfig,
  provider: string,
  model: string,
): Record<string, unknown> | null {
  const candidates: MatchResult[] = [];

  for (const [cfgProvider, models] of Object.entries(cfg)) {
    for (const [cfgModel, params] of Object.entries(models)) {
      const spec = matchSpecificity(cfgProvider, cfgModel, provider, model);
      if (spec >= 0) {
        candidates.push({ params, specificity: spec });
      }
    }
  }

  if (candidates.length === 0) return null;

  // Sort by specificity ascending (least specific first) then merge
  candidates.sort((a, b) => a.specificity - b.specificity);

  const merged: Record<string, unknown> = {};
  for (const candidate of candidates) {
    Object.assign(merged, candidate.params);
  }

  return merged;
}

/**
 * Returns specificity level (0-3) if the config entry matches,
 * or -1 if it doesn't match at all.
 */
function matchSpecificity(
  cfgProvider: string,
  cfgModel: string,
  targetProvider: string,
  targetModel: string,
): number {
  const providerMatch = matches(cfgProvider, targetProvider);
  const modelMatch = matches(cfgModel, targetModel);

  if (!providerMatch || !modelMatch) return -1;

  // Calculate specificity: count non-wildcard segments
  let spec = 0;
  if (cfgProvider !== "*") spec += 2; // exact provider
  if (cfgModel !== "*") spec += 1; // exact model
  return spec;
}

/**
 * Check if a config key matches a target value, supporting "*" wildcard.
 */
function matches(pattern: string, target: string): boolean {
  return pattern === "*" || pattern === target;
}

// ── Payload injection ────────────────────────────────────────

/**
 * Filter params to only supported keys, and inject them into the payload.
 * Returns a new payload object with overrides applied.
 */
function injectParams(
  payload: Record<string, unknown>,
  params: Record<string, unknown>,
): Record<string, unknown> {
  const overrides: Record<string, unknown> = {};

  for (const [key, value] of Object.entries(params)) {
    if (SUPPORTED_KEYS.has(key)) {
      overrides[key] = value;
    }
  }

  if (Object.keys(overrides).length === 0) return payload;

  return { ...payload, ...overrides };
}

// ── Value parsing ────────────────────────────────────────────

/**
 * Parse a string value into the appropriate type for the given key.
 * Numbers stay numbers, booleans stay booleans, known int fields are coerced.
 */
function parseValue(key: string, raw: string): unknown {
  // Known integer fields
  const intKeys = new Set(["top_k", "max_tokens"]);
  if (intKeys.has(key)) {
    const n = Number(raw);
    if (Number.isFinite(n) && Number.isInteger(n)) return n;
    return raw; // fallback to string
  }

  // Known float fields
  const floatKeys = new Set([
    "temperature",
    "top_p",
    "frequency_penalty",
    "presence_penalty",
    "repetition_penalty",
  ]);
  if (floatKeys.has(key)) {
    const n = Number(raw);
    if (Number.isFinite(n)) return n;
    return raw;
  }

  // Arrays (e.g. stop sequences): try JSON parse, else treat as string
  if (key === "stop") {
    try {
      return JSON.parse(raw);
    } catch {
      return raw;
    }
  }

  return raw;
}

// ── Override display ─────────────────────────────────────────

/**
 * Format a combined params object for display in the TUI.
 */
function formatOverrides(params: Record<string, unknown>): string {
  const lines = Object.entries(params)
    .filter(([, v]) => v !== undefined && v !== null)
    .map(([k, v]) => {
      const val = Array.isArray(v) ? JSON.stringify(v) : String(v);
      return `  ${k} = ${val}`;
    });
  return lines.length > 0 ? lines.join("\n") : "  (none)";
}

// ── Extension factory ────────────────────────────────────────

export default function (pi: ExtensionAPI): void {
  // ── Initialize config on any session start ──
  pi.on("session_start", async (_event, ctx) => {
    config = loadConfig(ctx.cwd);
  });

  // ── Track provider/model changes ──
  pi.on("model_select", async (event, _ctx) => {
    current.provider = event.model.provider;
    current.model = event.model.id;

    if (config) {
      const match = findBestMatch(config, current.provider, current.model);
      if (match) {
        const applied = Object.keys(match).filter((k) => SUPPORTED_KEYS.has(k));
        if (applied.length > 0) {
          console.log(
            `[pi-params] ${current.provider}/${current.model} -> applied: ${applied.join(", ")}`,
          );
        }
      }
    }
  });

  // ── Register /params slash commands ──
  pi.registerCommand("params", {
    description: "View or set model parameter overrides",
    usage: "/params [set|unset|preset|reset] [args...]",
    handler: async (args, ctx) => {
      const parts = args.trim().split(/\s+/);
      const subcommand = parts[0]?.toLowerCase();

      // ── /params (show) ──
      if (!subcommand || subcommand === "show") {
        const prov = current.provider ?? "(unknown)";
        const mdl = current.model ?? "(unknown)";
        let msg = `[pi-params] Current model: ${prov}/${mdl}\n\n`;

        // TOML config match
        if (config && current.provider && current.model) {
          const match = findBestMatch(config, current.provider, current.model);
          if (match) {
            const filtered = Object.fromEntries(
              Object.entries(match).filter(([k]) => SUPPORTED_KEYS.has(k)),
            );
            msg += `Config overrides:\n${formatOverrides(filtered)}\n`;
          } else {
            msg += "Config overrides: (none)\n";
          }
        } else {
          msg += "Config overrides: (not loaded)\n";
        }

        // Session overrides
        msg += `\nSession overrides:\n${formatOverrides(sessionOverrides)}`;

        // Presets hint
        msg +=
          "\n\nAvailable presets: " +
          Object.keys(PRESETS).join(", ") +
          "\nUsage: /params preset <name>";

        ctx.ui.notify(msg, "info");
        return;
      }

      // ── /params set <key> <value> ──
      if (subcommand === "set") {
        if (parts.length < 3) {
          ctx.ui.notify(
            "[pi-params] Usage: /params set <key> <value>\n" +
              "Example: /params set temperature 0.7",
            "warn",
          );
          return;
        }

        const key = parts[1];
        const rawValue = parts.slice(2).join(" ");

        if (!SUPPORTED_KEYS.has(key)) {
          ctx.ui.notify(
            `[pi-params] Unsupported key: "${key}". Supported: ${[...SUPPORTED_KEYS].join(", ")}`,
            "warn",
          );
          return;
        }

        const value = parseValue(key, rawValue);
        sessionOverrides[key] = value;

        ctx.ui.notify(
          `[pi-params] Set ${key} = ${Array.isArray(value) ? JSON.stringify(value) : String(value)}`,
          "info",
        );
        return;
      }

      // ── /params unset <key> ──
      if (subcommand === "unset") {
        if (parts.length < 2) {
          ctx.ui.notify("[pi-params] Usage: /params unset <key>", "warn");
          return;
        }

        const key = parts[1];
        if (key in sessionOverrides) {
          delete sessionOverrides[key];
          ctx.ui.notify(`[pi-params] Removed ${key} from session overrides`, "info");
        } else {
          ctx.ui.notify(`[pi-params] ${key} is not set in session overrides`, "warn");
        }
        return;
      }

      // ── /params preset <name> ──
      if (subcommand === "preset") {
        if (parts.length < 2) {
          ctx.ui.notify(
            "[pi-params] Usage: /params preset <name>\nPresets: " +
              Object.keys(PRESETS).join(", "),
            "warn",
          );
          return;
        }

        const name = parts[1].toLowerCase();
        const preset = PRESETS[name];
        if (!preset) {
          ctx.ui.notify(
            `[pi-params] Unknown preset: "${name}". Available: ${Object.keys(PRESETS).join(", ")}`,
            "warn",
          );
          return;
        }

        // Merge preset into session overrides (preset values take precedence)
        Object.assign(sessionOverrides, preset);
        ctx.ui.notify(
          `[pi-params] Applied preset "${name}": ${JSON.stringify(preset)}`,
          "info",
        );
        return;
      }

      // ── /params reset ──
      if (subcommand === "reset") {
        const count = Object.keys(sessionOverrides).length;
        sessionOverrides = {};
        ctx.ui.notify(
          `[pi-params] Cleared ${count} session override(s)`,
          "info",
        );
        return;
      }

      // Unknown subcommand
      ctx.ui.notify(
        "[pi-params] Unknown subcommand. Usage:\n" +
          "  /params                    — show overrides\n" +
          "  /params set <key> <value>  — set override\n" +
          "  /params unset <key>        — remove override\n" +
          "  /params preset <name>      — apply preset\n" +
          "  /params reset              — clear all overrides",
        "warn",
      );
    },
  });

  // ── Inject parameter overrides into provider requests ──
  pi.on("before_provider_request", (event, _ctx) => {
    // Skip if config not loaded or no model tracked
    const prov = current.provider;
    const mdl = current.model;
    if (!config || !prov || !mdl) return;

    // Find best matching config entry (merge from TOML)
    const match = findBestMatch(config, prov, mdl);
    if (!match && Object.keys(sessionOverrides).length === 0) return;

    // Validate payload is an object we can extend
    const payload = event.payload;
    if (typeof payload !== "object" || payload === null) {
      console.warn(
        `[pi-params] unexpected payload type: ${typeof payload}, skipping injection`,
      );
      return;
    }

    // Merge: TOML config first, then session overrides (highest priority)
    let mergedParams: Record<string, unknown>;
    if (match) {
      mergedParams = { ...match };
    } else {
      mergedParams = {};
    }
    Object.assign(mergedParams, sessionOverrides);

    return injectParams(payload as Record<string, unknown>, mergedParams);
  });
}