import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { dirname, join, resolve } from "node:path";
import { fileURLToPath } from "node:url";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { FALLBACK_MODELS } from "./models.ts";

const PROVIDER_NAME = "chutes-ai";
const DEFAULT_BASE_URL = "https://llm.chutes.ai/v1";
const API_KEY_ENV = "CHUTES_API_KEY";
const CONFIG_FILE = "pi-chutes.json";
const EXTENSION_DIR = dirname(fileURLToPath(import.meta.url));
const MODELS_CACHE_PATH = join(homedir(), ".pi", "agent", "extensions", "pi-chutes", "models.json");

type PiChutesConfig = {
  baseUrl: string;
  autoRefreshModels: boolean;
  recommendedModels: string[];
  hideNonRecommendedModels: boolean;
};

const DEFAULT_CONFIG: PiChutesConfig = {
  baseUrl: DEFAULT_BASE_URL,
  autoRefreshModels: true,
  recommendedModels: [],
  hideNonRecommendedModels: false,
};

type PiModel = {
  id: string;
  name: string;
  reasoning: boolean;
  input: ("text" | "image")[];
  cost: {
    input: number;
    output: number;
    cacheRead: number;
    cacheWrite: number;
  };
  contextWindow: number;
  maxTokens: number;
  compat?: {
    supportsDeveloperRole?: boolean;
    supportsReasoningEffort?: boolean;
    supportsUsageInStreaming?: boolean;
    maxTokensField?: "max_completion_tokens" | "max_tokens";
    thinkingFormat?: "openai" | "zai" | "qwen" | "qwen-chat-template";
  };
};

type ChutesModel = {
  id: string;
  pricing?: {
    prompt?: number;
    completion?: number;
    input_cache_read?: number;
  };
  input_modalities?: string[];
  supported_features?: string[];
  context_length?: number;
  max_model_len?: number;
  max_output_length?: number;
};

function loadCachedModels(): PiModel[] | null {
  if (!existsSync(MODELS_CACHE_PATH)) return null;

  try {
    const parsed = JSON.parse(readFileSync(MODELS_CACHE_PATH, "utf8")) as { models?: PiModel[] };
    return Array.isArray(parsed.models) ? parsed.models.map(withChutesCompat) : null;
  } catch {
    return null;
  }
}

function saveCachedModels(models: PiModel[]) {
  mkdirSync(dirname(MODELS_CACHE_PATH), { recursive: true });
  writeFileSync(
    MODELS_CACHE_PATH,
    JSON.stringify({ updatedAt: new Date().toISOString(), models: models.map(withChutesCompat) }, null, 2),
  );
  chmodSync(MODELS_CACHE_PATH, 0o600);
}

function compareModels(previous: PiModel[], next: PiModel[]) {
  const previousMap = new Map(previous.map((model) => [model.id, model]));
  const nextMap = new Map(next.map((model) => [model.id, model]));

  const added = next.filter((model) => !previousMap.has(model.id));
  const removed = previous.filter((model) => !nextMap.has(model.id));
  const updated = next.filter((model) => {
    const before = previousMap.get(model.id);
    return before ? JSON.stringify(before) !== JSON.stringify(model) : false;
  });

  return { added, removed, updated };
}

function isStringArray(value: unknown): value is string[] {
  return Array.isArray(value) && value.every((item) => typeof item === "string");
}

function loadConfig(): PiChutesConfig {
  const candidates = [
    resolve(process.cwd(), CONFIG_FILE),
    join(dirname(EXTENSION_DIR), CONFIG_FILE),
  ];

  for (const path of candidates) {
    if (!existsSync(path)) continue;

    try {
      const parsed = JSON.parse(readFileSync(path, "utf8")) as Record<string, unknown>;
      return {
        baseUrl: typeof parsed.baseUrl === "string" && parsed.baseUrl.trim() ? parsed.baseUrl : DEFAULT_CONFIG.baseUrl,
        autoRefreshModels:
          typeof parsed.autoRefreshModels === "boolean" ? parsed.autoRefreshModels : DEFAULT_CONFIG.autoRefreshModels,
        recommendedModels: isStringArray(parsed.recommendedModels) ? parsed.recommendedModels : DEFAULT_CONFIG.recommendedModels,
        hideNonRecommendedModels:
          typeof parsed.hideNonRecommendedModels === "boolean"
            ? parsed.hideNonRecommendedModels
            : DEFAULT_CONFIG.hideNonRecommendedModels,
      };
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      throw new Error(`Failed to parse ${path}: ${message}`);
    }
  }

  return DEFAULT_CONFIG;
}

function applyConfigToModels(models: PiModel[], config: PiChutesConfig): PiModel[] {
  const normalized = models.map(withChutesCompat);
  const recommended = new Set(config.recommendedModels);

  if (config.hideNonRecommendedModels && recommended.size > 0) {
    const filtered = normalized.filter((model) => recommended.has(model.id));
    return filtered.length > 0 ? filtered : normalized;
  }

  if (recommended.size === 0) return normalized;

  return [...normalized].sort((a, b) => {
    const aRecommended = recommended.has(a.id) ? 1 : 0;
    const bRecommended = recommended.has(b.id) ? 1 : 0;
    return bRecommended - aRecommended || a.id.localeCompare(b.id);
  });
}

function withChutesCompat(model: PiModel): PiModel {
  return {
    ...model,
    compat: {
      // Conservative defaults for OpenAI-compatible gateways.
      supportsDeveloperRole: false,
      supportsReasoningEffort: false,
      supportsUsageInStreaming: false,
      maxTokensField: "max_tokens",
      // Chutes' hosted models consistently accepted chat_template_kwargs.enable_thinking
      // in direct probes across multiple reasoning-capable backends (Qwen, DeepSeek,
      // GLM, Hermes). Using qwen-chat-template here aligns pi's thinking control with
      // the provider behavior and avoids leaked <think> tags / reasoning-only output
      // when thinking is meant to be off.
      ...(model.reasoning ? { thinkingFormat: "qwen-chat-template" as const } : {}),
      ...model.compat,
    },
  };
}

function toPiModel(model: ChutesModel): PiModel {
  const inputModalities = model.input_modalities ?? [];
  const supportsImage = inputModalities.includes("image");
  const supportedFeatures = model.supported_features ?? [];
  const pricing = model.pricing ?? {};

  return withChutesCompat({
    id: model.id,
    name: model.id,
    reasoning: supportedFeatures.includes("reasoning"),
    input: supportsImage ? ["text", "image"] : ["text"],
    cost: {
      input: pricing.prompt ?? 0,
      output: pricing.completion ?? 0,
      cacheRead: pricing.input_cache_read ?? 0,
      cacheWrite: 0,
    },
    contextWindow: model.context_length ?? model.max_model_len ?? 128000,
    maxTokens: model.max_output_length ?? 16384,
  });
}

async function fetchLiveModels(config: PiChutesConfig): Promise<PiModel[]> {
  const modelsUrl = `${config.baseUrl}/models`;
  const response = await fetch(modelsUrl, {
    headers: {
      Accept: "application/json",
    },
  });

  if (!response.ok) {
    throw new Error(`Failed to fetch ${modelsUrl}: ${response.status} ${response.statusText}`);
  }

  const payload = (await response.json()) as { data?: ChutesModel[] };
  const models = (payload.data ?? [])
    .filter((model): model is ChutesModel => Boolean(model?.id))
    .map(toPiModel);

  if (models.length === 0) {
    throw new Error("Chutes returned an empty model list");
  }

  return applyConfigToModels(models, config);
}

function registerProvider(pi: ExtensionAPI, models: PiModel[], config: PiChutesConfig) {
  pi.registerProvider(PROVIDER_NAME, {
    baseUrl: config.baseUrl,
    apiKey: API_KEY_ENV,
    authHeader: true,
    api: "openai-completions",
    models: applyConfigToModels(models, config),
    oauth: {
      name: "Chutes.ai",
      async login(callbacks: { onPrompt(params: { message: string }): Promise<string> }) {
        const apiKey = await callbacks.onPrompt({ message: "Enter your Chutes.ai API key:" });
        const token = apiKey.trim();
        if (!token) throw new Error("API key is required");
        return {
          access: token,
          refresh: token,
          expires: Number.MAX_SAFE_INTEGER,
        };
      },
      async refreshToken(credentials: { access: string; refresh: string; expires: number }) {
        return credentials;
      },
      getApiKey(credentials: { access: string }) {
        return credentials.access;
      },
    },
  });
}

export default function (pi: ExtensionAPI) {
  let config = loadConfig();
  let currentModels = applyConfigToModels(loadCachedModels() ?? [...FALLBACK_MODELS], config);

  registerProvider(pi, currentModels, config);

  const refresh = async (notify?: (message: string, level?: "info" | "success" | "warning" | "error") => void) => {
    try {
      config = loadConfig();
      const cachedOrFallback = applyConfigToModels(loadCachedModels() ?? [...FALLBACK_MODELS], config);
      currentModels = cachedOrFallback;
      registerProvider(pi, currentModels, config);

      if (!config.autoRefreshModels) {
        notify?.(
          `Chutes provider using cached/fallback catalog (${currentModels.length} models) from ${MODELS_CACHE_PATH}.`,
          "info",
        );
        return;
      }

      const liveModels = await fetchLiveModels(config);
      const { added, removed, updated } = compareModels(currentModels, liveModels);
      currentModels = liveModels;
      registerProvider(pi, currentModels, config);
      saveCachedModels(currentModels);

      const changes = [
        added.length > 0 ? `+${added.length} added` : null,
        removed.length > 0 ? `-${removed.length} removed` : null,
        updated.length > 0 ? `~${updated.length} updated` : null,
      ].filter(Boolean);

      notify?.(
        changes.length > 0
          ? `Chutes provider refreshed (${currentModels.length} models; ${changes.join(", ")}).`
          : `Chutes provider refreshed (${currentModels.length} models; no changes).`,
        "success",
      );
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      notify?.(`Chutes refresh failed; keeping cached/fallback catalog. ${message}`, "warning");
    }
  };

  pi.on("session_start", async (_event, ctx) => {
    await refresh((message, level = "info") => {
      if (ctx.hasUI) ctx.ui.notify(message, level);
    });
  });

  pi.registerCommand("chutes-refresh-models", {
    description: "Refresh the chutes.ai model catalog using current pi-chutes.json settings",
    handler: async (_args, ctx) => {
      await refresh((message, level = "info") => {
        if (ctx.hasUI) ctx.ui.notify(message, level);
      });
    },
  });
}