import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { AuthStorage } from "@earendil-works/pi-coding-agent";
import { randomUUID } from "node:crypto";

const CODEX_PROVIDER = "openai-codex";
const CHATGPT_BACKEND = "https://chatgpt.com/backend-api";

type RateLimitWindow = {
  used_percent?: number;
  limit_window_seconds?: number;
  reset_after_seconds?: number;
  reset_at?: number;
};

type ResetCredit = {
  id: string;
  reset_type?: string;
  status?: string;
  granted_at?: string;
  expires_at?: string | null;
  title?: string | null;
  description?: string | null;
};

type UsagePayload = {
  plan_type?: string;
  rate_limit?: {
    primary_window?: RateLimitWindow;
    secondary_window?: RateLimitWindow;
    limit_reached?: boolean;
  };
  rate_limit_reached_type?: { type?: string } | string | null;
  rate_limit_reset_credits?: { available_count?: number };
  additional_rate_limits?: Array<{
    limit_name?: string;
    metered_feature?: string;
    rate_limit?: { primary_window?: RateLimitWindow };
  }>;
};

type ResetCreditsPayload = {
  available_count: number;
  credits?: ResetCredit[];
};

type ConsumeOutcome = "reset" | "alreadyRedeemed" | "nothingToReset" | "noCredit";

type CodexAuth = {
  accessToken: string;
  accountId?: string;
};

type LimitsSnapshot = {
  usage: UsagePayload;
  resetCredits: ResetCreditsPayload;
};

function parseSubcommand(args: string): "view" | "reset" | null {
  const [subcommand] = args.trim().split(/\s+/);
  if (!subcommand || subcommand === "view") {
    return "view";
  }
  if (subcommand === "reset") {
    return "reset";
  }
  return null;
}

function formatWindow(seconds: number | undefined): string {
  if (seconds === undefined) {
    return "?";
  }
  if (seconds >= 86_400) {
    return `${Math.round(seconds / 86_400)}d`;
  }
  if (seconds >= 3_600) {
    return `${Math.round(seconds / 3_600)}h`;
  }
  return `${Math.round(seconds / 60)}m`;
}

function formatResetAt(timestamp: number | undefined): string {
  if (timestamp === undefined) {
    return "unknown";
  }
  return new Date(timestamp * 1_000).toLocaleString();
}

function formatIsoTimestamp(value: string | null | undefined): string {
  if (!value) {
    return "none";
  }
  const parsed = Date.parse(value);
  if (Number.isNaN(parsed)) {
    return value;
  }
  return new Date(parsed).toLocaleString();
}

function formatWindowLine(label: string, window: RateLimitWindow | undefined): string {
  if (!window) {
    return `${label}: unavailable`;
  }
  const used = window.used_percent ?? "?";
  const duration = formatWindow(window.limit_window_seconds);
  const resetsAt = formatResetAt(window.reset_at);
  return `${label}: ${used}% used (${duration} window, resets ${resetsAt})`;
}

function formatReachedType(value: UsagePayload["rate_limit_reached_type"]): string | null {
  if (!value) {
    return null;
  }
  if (typeof value === "string") {
    return value;
  }
  return value.type ?? null;
}

function formatLimitsReport(snapshot: LimitsSnapshot): string {
  const { usage, resetCredits } = snapshot;
  const lines = ["Codex rate limits"];

  if (usage.plan_type) {
    lines.push(`Plan: ${usage.plan_type}`);
  }

  lines.push(formatWindowLine("Primary", usage.rate_limit?.primary_window));
  lines.push(formatWindowLine("Secondary", usage.rate_limit?.secondary_window));

  const reachedType = formatReachedType(usage.rate_limit_reached_type);
  if (reachedType) {
    lines.push(`Limit state: ${reachedType}`);
  }

  for (const extra of usage.additional_rate_limits ?? []) {
    const label = extra.limit_name ?? extra.metered_feature ?? "Additional";
    lines.push(formatWindowLine(label, extra.rate_limit?.primary_window));
  }

  const availableCount = resetCredits.available_count;
  lines.push(`Reset credits available: ${availableCount}`);

  const credits = resetCredits.credits ?? [];
  for (const credit of credits) {
    const title = credit.title ?? credit.id;
    const expires = formatIsoTimestamp(credit.expires_at);
    lines.push(`  - ${title} (${credit.status ?? "unknown"}, expires ${expires})`);
  }

  if (availableCount > credits.length && credits.length > 0) {
    lines.push(`  (${availableCount - credits.length} more not listed by backend)`);
  }

  return lines.join("\n");
}

function backendUrls(baseUrl: string) {
  const root = baseUrl.replace(/\/$/, "");
  const prefix = root.includes("/backend-api") ? `${root}/wham` : `${root}/api/codex`;
  return {
    usage: `${prefix}/usage`,
    resetCredits: `${prefix}/rate-limit-reset-credits`,
    consumeResetCredit: `${prefix}/rate-limit-reset-credits/consume`,
  };
}

async function loadCodexAuth(): Promise<CodexAuth> {
  const authStorage = AuthStorage.create();
  const credential = authStorage.get(CODEX_PROVIDER);
  if (!credential || credential.type !== "oauth") {
    throw new Error("ChatGPT OAuth is required. Run /login and choose openai-codex.");
  }

  const accessToken = await authStorage.getApiKey(CODEX_PROVIDER);
  if (!accessToken) {
    throw new Error("Could not load a valid openai-codex access token. Run /login.");
  }

  const accountId =
    typeof credential.accountId === "string" && credential.accountId.length > 0
      ? credential.accountId
      : undefined;

  return { accessToken, accountId };
}

async function codexFetch<T>(auth: CodexAuth, url: string, init?: RequestInit): Promise<T> {
  const headers = new Headers(init?.headers);
  headers.set("Authorization", `Bearer ${auth.accessToken}`);
  headers.set("User-Agent", "pi-codex-limits-extension");
  headers.set("Accept", "application/json");
  if (auth.accountId) {
    headers.set("ChatGPT-Account-ID", auth.accountId);
  }
  if (init?.body) {
    headers.set("Content-Type", "application/json");
  }

  const response = await fetch(url, { ...init, headers });
  const body = await response.text();
  if (!response.ok) {
    throw new Error(`${init?.method ?? "GET"} ${url} failed (${response.status}): ${body}`);
  }
  return JSON.parse(body) as T;
}

async function fetchLimitsSnapshot(auth: CodexAuth): Promise<LimitsSnapshot> {
  const urls = backendUrls(CHATGPT_BACKEND);
  const usage = await codexFetch<UsagePayload>(auth, urls.usage);

  let resetCredits: ResetCreditsPayload = {
    available_count: usage.rate_limit_reset_credits?.available_count ?? 0,
    credits: [],
  };

  try {
    resetCredits = await codexFetch<ResetCreditsPayload>(auth, urls.resetCredits);
  } catch (error) {
    if (resetCredits.available_count === 0) {
      throw error;
    }
  }

  if (resetCredits.available_count === 0 && usage.rate_limit_reset_credits?.available_count) {
    resetCredits.available_count = usage.rate_limit_reset_credits.available_count;
  }

  return { usage, resetCredits };
}

function availableCredits(snapshot: LimitsSnapshot): ResetCredit[] {
  const credits = snapshot.resetCredits.credits ?? [];
  const available = credits.filter((credit) => credit.status === "available");
  return available.length > 0 ? available : credits;
}

function outcomeMessage(outcome: ConsumeOutcome, windowsReset: number): string {
  switch (outcome) {
    case "reset":
      return `Reset credit consumed${windowsReset > 0 ? ` (${windowsReset} window${windowsReset === 1 ? "" : "s"} reset)` : ""}.`;
    case "alreadyRedeemed":
      return "This reset was already redeemed. Refreshing limits.";
    case "nothingToReset":
      return "No eligible rate-limit window to reset.";
    case "noCredit":
      return "No reset credits available.";
  }
}

function mapConsumeCode(code: string): ConsumeOutcome {
  switch (code) {
    case "reset":
      return "reset";
    case "already_redeemed":
      return "alreadyRedeemed";
    case "nothing_to_reset":
      return "nothingToReset";
    case "no_credit":
      return "noCredit";
    default:
      throw new Error(`Unexpected reset response code: ${code}`);
  }
}

async function consumeResetCredit(
  auth: CodexAuth,
  creditId?: string,
): Promise<{ outcome: ConsumeOutcome; windowsReset: number }> {
  const urls = backendUrls(CHATGPT_BACKEND);
  const payload: { redeem_request_id: string; credit_id?: string } = {
    redeem_request_id: randomUUID(),
  };
  if (creditId) {
    payload.credit_id = creditId;
  }

  const response = await codexFetch<{ code: string; windows_reset?: number }>(
    auth,
    urls.consumeResetCredit,
    {
      method: "POST",
      body: JSON.stringify(payload),
    },
  );

  return {
    outcome: mapConsumeCode(response.code),
    windowsReset: response.windows_reset ?? 0,
  };
}

async function handleView(ctx: ExtensionCommandContext): Promise<void> {
  const auth = await loadCodexAuth();
  const snapshot = await fetchLimitsSnapshot(auth);
  ctx.ui.notify(formatLimitsReport(snapshot), "info");
}

async function handleReset(ctx: ExtensionCommandContext): Promise<void> {
  const auth = await loadCodexAuth();
  const snapshot = await fetchLimitsSnapshot(auth);
  const credits = availableCredits(snapshot);

  if (snapshot.resetCredits.available_count <= 0 || credits.length === 0) {
    ctx.ui.notify("No reset credits available.", "warning");
    return;
  }

  let creditId: string | undefined;
  if (credits.length === 1) {
    creditId = credits[0]?.id;
  } else {
    const selected = await ctx.ui.select(
      "Choose a reset credit",
      credits.map((credit) => {
        const title = credit.title ?? credit.id;
        const expires = credit.expires_at ? `, expires ${formatIsoTimestamp(credit.expires_at)}` : "";
        return `${title}${expires}`;
      }),
    );
    if (!selected) {
      ctx.ui.notify("Reset cancelled.", "info");
      return;
    }
    const index = credits.findIndex((credit) => {
      const title = credit.title ?? credit.id;
      const expires = credit.expires_at ? `, expires ${formatIsoTimestamp(credit.expires_at)}` : "";
      return `${title}${expires}` === selected;
    });
    creditId = credits[index]?.id;
  }

  const credit = credits.find((entry) => entry.id === creditId);
  const label = credit?.title ?? creditId ?? "next available credit";
  const confirmed = await ctx.ui.confirm(
    "Consume reset credit?",
    `Redeem "${label}" and reset Codex rate limits?`,
  );
  if (!confirmed) {
    ctx.ui.notify("Reset cancelled.", "info");
    return;
  }

  const { outcome, windowsReset } = await consumeResetCredit(auth, creditId);
  const refreshed = await fetchLimitsSnapshot(auth);
  const message = `${outcomeMessage(outcome, windowsReset)}\n\n${formatLimitsReport(refreshed)}`;
  const level = outcome === "reset" || outcome === "alreadyRedeemed" ? "info" : "warning";
  ctx.ui.notify(message, level);
}

export default function (pi: ExtensionAPI) {
  pi.registerCommand("codex-limits", {
    description: "View or reset Codex rate limits (view|reset)",
    getArgumentCompletions: (prefix) => {
      const options = ["view", "reset"];
      const filtered = options.filter((option) => option.startsWith(prefix));
      return filtered.length > 0 ? filtered.map((value) => ({ value, label: value })) : null;
    },
    handler: async (args, ctx) => {
      const subcommand = parseSubcommand(args);
      if (!subcommand) {
        ctx.ui.notify("Usage: /codex-limits [view|reset]", "error");
        return;
      }

      ctx.ui.setStatus("codex-limits", `${subcommand}...`);
      try {
        if (subcommand === "view") {
          await handleView(ctx);
        } else {
          await handleReset(ctx);
        }
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        ctx.ui.notify(message, "error");
      } finally {
        ctx.ui.setStatus("codex-limits", undefined);
      }
    },
  });
}