import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
  lazyStream,
  type Api,
  type Context,
  type Model,
  type Provider,
  type ProviderHeaders,
  type SimpleStreamOptions,
  type StreamOptions,
} from "@earendil-works/pi-ai";

const PROVIDER_ID = "github-copilot";
const AUTO_MODEL_ID = "auto";
const AUTO_PREFIX = "auto-";

function realModelId(id: string): string {
  return id.startsWith(AUTO_PREFIX) ? id.slice(AUTO_PREFIX.length) : id;
}
const COPILOT_HEADERS = {
  Accept: "application/json",
  "Content-Type": "application/json",
  "User-Agent": "GitHubCopilotChat/0.35.0",
  "Editor-Version": "vscode/1.107.0",
  "Editor-Plugin-Version": "copilot-chat/0.35.0",
  "Copilot-Integration-Id": "vscode-chat",
  "X-GitHub-Api-Version": "2026-06-01",
  "Openai-Intent": "conversation-edits",
} as const;

const DEFAULT_BASE_URL = "https://api.individual.githubcopilot.com";

interface AutoSession {
  availableModels: string[];
  sessionToken: string;
  expiresAt: number;
  interactionId: string;
  chosenModel?: string;
  reasoningBucket?: "low" | "medium" | "high";
}

interface SessionResponse {
  available_models?: unknown;
  selected_model?: unknown;
  session_token?: unknown;
  expires_at?: unknown;
}

interface IntentResponse {
  chosen_model?: unknown;
  candidate_models?: unknown;
  reasoning_bucket?: unknown;
}

// Same API routing rule Pi's model catalog applies to github-copilot models:
// Claude 4.x/5.x -> anthropic-messages; gpt-5*/oswe*/mai-* -> openai-responses;
// everything else (gpt-4.x, gemini, ...) -> openai-completions.
function apiForModel(id: string): Api {
  if (/^claude-(haiku|sonnet|opus)-[45]([.-]|$)/.test(id)) return "anthropic-messages";
  if (id.startsWith("gpt-5") || id.startsWith("oswe") || id.startsWith("mai-"))
    return "openai-responses";
  return "openai-completions";
}

function poolModel(id: string, name: string, api: Api = apiForModel(id)): Model<Api> {
  const anthropic = api === "anthropic-messages";
  const responses = api === "openai-responses";
  let compat: Record<string, unknown>;
  if (anthropic) {
    compat = { supportsEagerToolInputStreaming: false };
  } else if (responses) {
    compat = {
      supportsReasoningEffort: true,
      supportsStore: false,
      supportsStrictMode: true,
      sessionAffinityFormat: "openai",
    };
  } else {
    compat = {
      supportsStore: false,
      supportsDeveloperRole: false,
      supportsReasoningEffort: false,
    };
  }
  return {
    id,
    name,
    api,
    provider: PROVIDER_ID,
    baseUrl: DEFAULT_BASE_URL,
    reasoning: responses,
    input: ["text", "image"],
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: 400_000,
    maxTokens: anthropic ? 64_000 : 128_000,
    thinkingLevelMap: responses
      ? {
          off: "none",
          minimal: "low",
          low: "low",
          medium: "medium",
          high: "high",
          xhigh: null,
          max: null,
        }
      : undefined,
    compat,
  };
}

function latestUserPrompt(context: Context): { prompt: string; hasImage: boolean } {
  for (let index = context.messages.length - 1; index >= 0; index--) {
    const message = context.messages[index];
    if (message.role !== "user") continue;
    if (typeof message.content === "string") return { prompt: message.content, hasImage: false };

    let prompt = "";
    let hasImage = false;
    for (const part of message.content) {
      if (part.type === "text") prompt += `${prompt ? "\n" : ""}${part.text}`;
      if (part.type === "image") hasImage = true;
    }
    return { prompt, hasImage };
  }
  return { prompt: "", hasImage: false };
}

function mergeHeaders(
  base: ProviderHeaders | undefined,
  extra: Record<string, string>,
): ProviderHeaders {
  const merged: ProviderHeaders = { ...base };
  for (const [name, value] of Object.entries(extra)) {
    for (const existing of Object.keys(merged)) {
      if (existing.toLowerCase() === name.toLowerCase()) delete merged[existing];
    }
    merged[name] = value;
  }
  return merged;
}

async function fetchJson<T>(url: string, init: RequestInit, signal?: AbortSignal): Promise<T> {
  const timeout = AbortSignal.timeout(15_000);
  const combinedSignal = signal ? AbortSignal.any([signal, timeout]) : timeout;
  const response = await fetch(url, { ...init, signal: combinedSignal });
  if (!response.ok) {
    const body = await response.text();
    throw new Error(
      `Copilot Auto ${response.status} ${response.statusText}: ${body.slice(0, 500)}`,
    );
  }
  return (await response.json()) as T;
}

async function createAutoSession(
  baseUrl: string,
  apiKey: string,
  signal?: AbortSignal,
): Promise<AutoSession> {
  const interactionId = crypto.randomUUID();
  const response = await fetchJson<SessionResponse>(
    `${baseUrl}/models/session`,
    {
      method: "POST",
      headers: {
        ...COPILOT_HEADERS,
        Authorization: `Bearer ${apiKey}`,
        "X-Initiator": "user",
        "X-Interaction-Id": interactionId,
      },
      body: JSON.stringify({ auto_mode: { model_hints: [AUTO_MODEL_ID] } }),
    },
    signal,
  );

  const availableModels = Array.isArray(response.available_models)
    ? response.available_models.filter((value): value is string => typeof value === "string")
    : [];
  if (availableModels.length === 0 || typeof response.session_token !== "string") {
    throw new Error("Copilot Auto returned an invalid model session");
  }

  return {
    availableModels,
    sessionToken: response.session_token,
    expiresAt:
      typeof response.expires_at === "number"
        ? response.expires_at * 1000
        : Date.now() + 10 * 60_000,
    interactionId,
    chosenModel:
      typeof response.selected_model === "string" ? response.selected_model : availableModels[0],
  };
}

async function routePrompt(
  baseUrl: string,
  apiKey: string,
  state: AutoSession,
  context: Context,
  signal?: AbortSignal,
): Promise<void> {
  const { prompt, hasImage } = latestUserPrompt(context);
  const response = await fetchJson<IntentResponse>(
    `${baseUrl}/models/session/intent`,
    {
      method: "POST",
      headers: {
        ...COPILOT_HEADERS,
        Authorization: `Bearer ${apiKey}`,
        "Copilot-Session-Token": state.sessionToken,
        "X-Initiator": "user",
        "X-Interaction-Id": state.interactionId,
      },
      body: JSON.stringify({
        prompt,
        available_models: state.availableModels,
        has_image: hasImage,
      }),
    },
    signal,
  );

  const candidates = Array.isArray(response.candidate_models)
    ? response.candidate_models.filter((value): value is string => typeof value === "string")
    : [];
  const chosen = typeof response.chosen_model === "string" ? response.chosen_model : candidates[0];
  if (!chosen) throw new Error("Copilot Auto router did not choose a model");

  state.chosenModel = chosen;
  if (
    response.reasoning_bucket === "low" ||
    response.reasoning_bucket === "medium" ||
    response.reasoning_bucket === "high"
  ) {
    state.reasoningBucket = response.reasoning_bucket;
  } else {
    state.reasoningBucket = undefined;
  }
}

async function fetchAutoPool(baseUrl: string, apiKey: string): Promise<string[]> {
  try {
    const session = await createAutoSession(baseUrl, apiKey);
    return session.availableModels;
  } catch {
    return [];
  }
}

function wrapProvider(base: Provider, pool: string[]): Provider {
  const baseById = new Map(base.getModels().map((entry) => [entry.id, entry]));

  // Prefer the real catalog definition (correct api/compat/context window) and
  // fall back to a heuristic only for Auto-only models missing from the catalog.
  const templateFor = (realId: string, displayId: string): Model<Api> => {
    const known = baseById.get(realId);
    if (known) return { ...known, id: displayId, name: displayId };
    return poolModel(displayId, displayId, apiForModel(realId));
  };

  const routerModel = poolModel(AUTO_MODEL_ID, "Copilot Auto", "openai-responses");
  const poolIds = new Set(pool);
  const managedIds = new Set([AUTO_MODEL_ID, ...pool.map((id) => `${AUTO_PREFIX}${id}`)]);
  const poolModels = pool.map((id) => templateFor(id, `${AUTO_PREFIX}${id}`));
  const poolModelByRealId = new Map(pool.map((id, index) => [id, poolModels[index]]));
  const sessions = new Map<string, AutoSession>();

  async function prepare(
    requestModel: Model<Api>,
    context: Context,
    options: StreamOptions | undefined,
  ): Promise<{ model: Model<Api>; options: SimpleStreamOptions }> {
    const apiKey = options?.apiKey;
    if (!apiKey) throw new Error("GitHub Copilot authentication is unavailable");

    const forced = poolIds.has(realModelId(requestModel.id))
      ? realModelId(requestModel.id)
      : undefined;
    const baseUrl = requestModel.baseUrl ?? base.baseUrl ?? routerModel.baseUrl ?? DEFAULT_BASE_URL;

    const key = `${forced ?? "auto"}:${options?.sessionId ?? "default"}`;
    let state = sessions.get(key);
    if (!state || state.expiresAt <= Date.now() + 30_000) {
      state = await createAutoSession(baseUrl, apiKey, options?.signal);
      sessions.set(key, state);
    }

    if (forced) {
      state.chosenModel = forced;
      state.reasoningBucket = undefined;
    } else {
      const lastMessage = context.messages.at(-1);
      if (!state.chosenModel || lastMessage?.role === "user") {
        await routePrompt(baseUrl, apiKey, state, context, options?.signal);
      }
    }
    if (!state.chosenModel) throw new Error("Copilot Auto did not select a model");

    const template =
      poolModelByRealId.get(state.chosenModel) ?? templateFor(state.chosenModel, state.chosenModel);
    return {
      model: { ...template, id: state.chosenModel, name: state.chosenModel, baseUrl },
      options: {
        ...options,
        reasoning: state.reasoningBucket ?? (options as SimpleStreamOptions | undefined)?.reasoning,
        headers: mergeHeaders(options?.headers, {
          ...COPILOT_HEADERS,
          "Copilot-Session-Token": state.sessionToken,
          "X-Interaction-Id": state.interactionId,
        }),
      },
    };
  }

  const streamAuto = (requestModel: Model<Api>, context: Context, options?: StreamOptions) =>
    lazyStream(requestModel, async () => {
      const routed = await prepare(requestModel, context, options);
      return base.streamSimple(routed.model, context, routed.options);
    });

  return {
    id: base.id,
    name: base.name,
    baseUrl: base.baseUrl,
    headers: base.headers,
    auth: base.auth,
    getModels: () => {
      const models = base.getModels().filter((entry) => !managedIds.has(entry.id));
      return [routerModel, ...poolModels, ...models];
    },
    refreshModels: base.refreshModels ? (context) => base.refreshModels!(context) : undefined,
    filterModels: (models, credential) => {
      const remaining = models.filter((entry) => !managedIds.has(entry.id));
      const filtered = base.filterModels?.(remaining, credential) ?? remaining;
      return [routerModel, ...poolModels, ...filtered];
    },
    stream: (requestModel, context, options) =>
      managedIds.has(requestModel.id)
        ? streamAuto(requestModel, context, options)
        : base.stream(requestModel, context, options),
    streamSimple: (requestModel, context, options) =>
      managedIds.has(requestModel.id)
        ? streamAuto(requestModel, context, options)
        : base.streamSimple(requestModel, context, options),
  };
}

export default function githubCopilotAuto(pi: ExtensionAPI) {
  pi.on("session_start", (_event, ctx) => {
    const base = ctx.modelRegistry.getProvider(PROVIDER_ID);
    if (!base || base.getModels().some((entry) => entry.id === AUTO_MODEL_ID)) return;

    // Register immediately with the router model so startup is not blocked.
    // The per-account pool is fetched in the background and re-registered when ready.
    pi.registerProvider(wrapProvider(base, []));

    void (async () => {
      try {
        const resolved = await ctx.modelRegistry.getProviderAuth(PROVIDER_ID);
        const apiKey = resolved?.auth.apiKey;
        const baseUrl = resolved?.auth.baseUrl ?? base.baseUrl ?? DEFAULT_BASE_URL;
        if (!apiKey) return;
        const pool = await fetchAutoPool(baseUrl, apiKey);
        if (pool.length > 0) pi.registerProvider(wrapProvider(base, pool));
      } catch {
        // Keep the router-only provider already registered.
      }
    })();
  });
}