import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { readFileSync, writeFileSync, existsSync, mkdirSync } from "node:fs";
import { join } from "node:path";
import os from "node:os";

interface Preset {
  params: Record<string, any>;
  description?: string;
}

interface State {
  models: Record<string, Record<string, any>>;
  presets: Record<string, Record<string, Preset>>; // modelId -> presetName -> Preset
  shorthands: Record<string, string>;
}

const DEFAULT_STATE: State = {
  models: {},
  presets: {},
  shorthands: {
    temp: "temperature",
    tp: "top_p",
    tk: "top_k",
    mp: "min_p",
    rp: "repeat_penalty",
    pp: "presence_penalty",
  },
};

function getGlobalStatePath() {
  return join(os.homedir(), ".pi", "agent", "generation-config.json");
}

function loadState(): State {
  const stateFile = getGlobalStatePath();
  if (!existsSync(stateFile)) return { ...DEFAULT_STATE };
  try {
    const data = JSON.parse(readFileSync(stateFile, "utf8"));
    return {
      models: data.models ?? DEFAULT_STATE.models,
      presets: (data.presets as Record<string, Record<string, Preset>>) ?? DEFAULT_STATE.presets,
      shorthands: { ...DEFAULT_STATE.shorthands, ...(data.shorthands || {}) },
    };
  } catch (e) {
    console.error("Failed to load generation-config state:", e);
    return { ...DEFAULT_STATE };
  }
}

function saveState(state: State) {
  try {
    const stateFile = getGlobalStatePath();
    const dir = join(os.homedir(), ".pi", "agent");
    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true });
    }
    writeFileSync(stateFile, JSON.stringify(state, null, 2));
  } catch (e) {
    console.error("Failed to save generation-config state:", e);
  }
}

function parseValue(val: string): any {
  const trimmed = val.trim();
  if (trimmed.toLowerCase() === "true") return true;
  if (trimmed.toLowerCase() === "false") return false;
  if (!isNaN(Number(trimmed)) && trimmed !== "") return Number(trimmed);
  return trimmed;
}

// Turn-level parameters (overrides)
let turnParams: Record<string, any> = {};

export default function (pi: ExtensionAPI) {
  let state: State = { ...DEFAULT_STATE };

  pi.on("session_start", (_event, ctx) => {
    state = loadState();
  });

  const resolveKey = (key: string) => state.shorthands[key.toLowerCase()] || key;

  pi.registerCommand("gconfig", {
    description: "Manage generation configuration",
    handler: async (args, ctx) => {
      if (!args) {
        const modelId = ctx.model.id;
        const settings = state.models[modelId] || {};
        const settingsList = Object.entries(settings)
          .map(([k, v]) => `${k}: ${v}`)
          .join("\n") || "None";

        const modelPresets = state.presets[modelId] ?? {};
        const presetsList = Object.entries(modelPresets)
          .map(([name, p]) => `${name}${p.description ? ` (${p.description})` : ""}`)
          .join("\n") || "None";

        ctx.ui.notify(
          `Active Model [${modelId}] Settings:\n${settingsList}\n\nAvailable Presets:\n${presetsList}\n\nUsage: /gconfig <set|set-model|params|reset|preset|shorthand|delete>`,
          "info"
        );
        return;
      }

      const parts = args.split(/\s+/);
      const subCommand = parts[0].toLowerCase();

      if (subCommand === "set") {
        const key = parts[1];
        const value = parts[2];
        const shorthand = parts[3];
        if (!key || value === undefined) {
          ctx.ui.notify("Usage: /gconfig set <key> <value> [shorthand]", "error");
          return;
        }
        const modelId = ctx.model.id;
        const finalKey = resolveKey(key);
        const finalVal = parseValue(value);

        if (!state.models[modelId]) state.models[modelId] = {};
        state.models[modelId][finalKey] = finalVal;

        if (shorthand) {
          state.shorthands[shorthand.toLowerCase()] = finalKey;
        }

        saveState(state);
        ctx.ui.notify(`Model [${modelId}] ${finalKey} set to ${finalVal}${shorthand ? ` (shorthand: ${shorthand})` : ""}`, "info");
      } 
      
      else if (subCommand === "set-model") {
        const key = parts[1];
        const value = parts[2];
        const modelId = parts[3];
        const shorthand = parts[4];
        if (!key || value === undefined || !modelId) {
          ctx.ui.notify("Usage: /gconfig set-model <key> <value> <model> [shorthand]", "error");
          return;
        }
        const finalKey = resolveKey(key);
        const finalVal = parseValue(value);

        if (!state.models[modelId]) state.models[modelId] = {};
        state.models[modelId][finalKey] = finalVal;

        if (shorthand) {
          state.shorthands[shorthand.toLowerCase()] = finalKey;
        }

        saveState(state);
        ctx.ui.notify(`Model [${modelId}] ${finalKey} set to ${finalVal}${shorthand ? ` (shorthand: ${shorthand})` : ""}`, "info");
      }

      else if (subCommand === "delete") {
        const type = parts[1]?.toLowerCase(); // 'param' or 'preset'
        const modelId = ctx.model.id;
        if (type === "param") {
          const key = parts[2];
          if (!key) {
            ctx.ui.notify("Usage: /gconfig delete param <key>", "error");
            return;
          }
          if (state.models[modelId] && state.models[modelId][resolveKey(key)] !== undefined) {
            delete state.models[modelId][resolveKey(key)];
            saveState(state);
            ctx.ui.notify(`Deleted parameter ${key} from ${modelId}`, "info");
          } else {
            ctx.ui.notify(`Parameter ${key} not found for ${modelId}`, "error");
          }
        }        
        else if (type === "preset") {
          const name = parts[2];
          if (!name) {
            ctx.ui.notify("Usage: /gconfig delete preset <name>", "error");
            return;
          }
          if (state.presets[modelId] && state.presets[modelId][name]) {
            delete state.presets[modelId][name];
            saveState(state);
            ctx.ui.notify(`Deleted preset '${name}' for ${modelId}`, "info");
          } else {
            ctx.ui.notify(`Preset '${name}' not found for ${modelId}`, "error");
          }
        }
        else {
          ctx.ui.notify("Usage: /gconfig delete <param|preset> <name/key>", "error");
        }
      }

      else if (subCommand === "params") {
        const modelId = ctx.model.id;
        const activeSettings = state.models[modelId] || {};
        
        const knownKeys = [...new Set(Object.values(state.shorthands))];
        const knownParamsList = knownKeys.map(k => {
          const shorthands = Object.entries(state.shorthands)
            .filter(([_, long]) => long === k)
            .map(([short]) => short)
            .join(", ");
          const val = activeSettings[k] !== undefined ? activeSettings[k] : "default";
          return `${k} [${shorthands}]: ${val}`;
        }).join("\n");

        let output = `Known Parameters (Active Model [${modelId}]:\n${knownParamsList}\n\n`;

        const modelIds = Object.keys(state.models);
        if (modelIds.length > 0) {
          output += `Model-Specific Configurations:\n`;
          const modelConfigs = modelIds.map(mId => {
            const settings = state.models[mId];
            const paramsList = Object.entries(settings).map(([k, v]) => {
              const shorthands = Object.entries(state.shorthands)
                .filter(([_, long]) => long === k)
                .map(([short]) => short)
                .join(", ");
              const shorthandStr = shorthands ? ` [${shorthands}]` : "";
              return `  ${k}${shorthandStr}: ${v}`;
            }).join("\n");
            return `${mId}:\n${paramsList}`;
          }).join("\n\n");
          output += modelConfigs;
        } else {
          output += `No model-specific overrides configured.`;
        }
        
        ctx.ui.notify(output, "info");
      }

      else if (subCommand === "reset") {
        const modelId = ctx.model.id;
        if (state.models[modelId]) {
          delete state.models[modelId];
          saveState(state);
          ctx.ui.notify(`Reset settings for ${modelId}`, "info");
        } else {
          ctx.ui.notify("No custom settings to reset", "info");
        }
      }

      else if (subCommand === "preset") {
        const action = parts[1]?.toLowerCase();
        const modelId = ctx.model.id;
        if (!state.presets[modelId]) state.presets[modelId] = {};
        if (action === "save") {
          const name = parts[2];
          const description = parts.slice(3).join(" ");
          if (!name) {
            ctx.ui.notify("Usage: /gconfig preset save <name> [description]", "error");
            return;
          }
          const currentSettings = { ...state.models[modelId] };
          if (Object.keys(currentSettings).length === 0) {
            ctx.ui.notify("No settings configured for current model to save as preset", "error");
            return;
          }
          state.presets[modelId][name] = { params: currentSettings, description };
          saveState(state);
          ctx.ui.notify(`Saved current settings to preset '${name}' for ${modelId}${description ? `: ${description}` : ""}`, "info");
        }
        else if (action === "load") {
          const name = parts[2];
          if (!name) {
            ctx.ui.notify("Usage: /gconfig preset load <name>", "error");
            return;
          }
          const preset = state.presets[modelId]?.[name];
          if (!preset) {
            ctx.ui.notify(`Preset '${name}' not found for ${modelId}`, "error");
            return;
          }
          if (!state.models[modelId]) state.models[modelId] = {};
          state.models[modelId] = { ...state.models[modelId], ...preset.params };
          saveState(state);
          ctx.ui.notify(`Applied preset '${name}' to ${modelId}`, "info");
        }
        else if (action === "list") {
          const entries = Object.entries(state.presets[modelId] ?? {});
          if (entries.length === 0) {
            ctx.ui.notify(`No presets saved for ${modelId}`, "info");
            return;
          }
          const list = entries.map(([name, p]) => `${name}${p.description ? ` (${p.description})` : ""}: ${JSON.stringify(p.params)}`).join("\n");
          ctx.ui.notify(`Presets for ${modelId}:\n${list}`, "info");
        }
        else {
          ctx.ui.notify("Usage: /gconfig preset <save|load|list> [name] [description]", "error");
        }
      }

      else if (subCommand === "shorthand") {
        const action = parts[1]?.toLowerCase();
        if (action === "add") {
          const [_, __, short, long] = parts;
          if (!short || !long) {
            ctx.ui.notify("Usage: /gconfig shorthand add <short> <long>", "error");
            return;
          }
          state.shorthands[short.toLowerCase()] = long.toLowerCase();
          saveState(state);
          ctx.ui.notify(`Added shorthand: ${short} -> ${long}`, "info");
        }
        else {
          ctx.ui.notify("Use /gconfig params to see shorthands. Use /gconfig shorthand add <short> <long> to add new ones.", "info");
        }
      }

      else {
        ctx.ui.notify("Unknown gconfig command. Use /gconfig for help.", "error");
      }
    },
  });

  pi.on("input", async (event, ctx) => {
    const tagRegex = /^\s*(\[[^\]]*\])/gi;
    let text = event.text;
    let hasTags = false;

    turnParams = {};

    while (true) {
      const tagMatch = text.match(tagRegex);
      if (!tagMatch) break;

      const tagContent = tagMatch[1];
      const innerContent = tagContent.slice(1, -1);
      const paramSegments = innerContent.split(",").map(s => s.trim()).filter(Boolean);

      for (const segment of paramSegments) {
        const colonIndex = segment.indexOf(":");
        if (colonIndex > 0) {
          const key = segment.slice(0, colonIndex).trim();
          const value = segment.slice(colonIndex + 1).trim();
          const resolvedKey = resolveKey(key);
          if (resolvedKey.toLowerCase() === "preset") {
            const presetName = value;
            const modelId = ctx.model.id;
            const preset = state.presets[modelId]?.[presetName];
            if (preset) {
              Object.assign(turnParams, preset.params);
            }
          } else {
            turnParams[resolvedKey] = parseValue(value);
          }
        }
      }

      text = text.replace(tagContent, "");
      hasTags = true;
    }

    if (hasTags) {
      return { action: "transform", text: text.trim() };
    }

    return { action: "continue" };
  });

  pi.on("turn_end", async () => {
    turnParams = {};
  });

  pi.on("before_provider_request", (event, ctx) => {
    const modelId = ctx.model.id;
    const modelSettings = state.models[modelId] || {};
    
    const mergedParams = {
      ...modelSettings,
      ...turnParams,
    };

    if (Object.keys(mergedParams).length === 0) return;

    return {
      ...event.payload,
      ...mergedParams,
    };
  });
}