import type {
  ExtensionAPI,
  ProviderModelConfig,
} from "@earendil-works/pi-coding-agent";
import type {
  OAuthCredentials,
  OAuthLoginCallbacks,
} from "@earendil-works/pi-ai";

const PROVIDER_NAME = "umans";
const PROVIDER_DISPLAY_NAME = "Umans";
// Host root only. The anthropic-messages adapter (the Anthropic SDK) appends
// `/v1/messages` itself.
const BASE_URL = "https://api.code.umans.ai";
const MODELS_INFO_URL = `${BASE_URL}/v1/models/info`;
const API_KEY_ENV_VAR = "UMANS_API_KEY";
const API_KEY_ENV_REF = `$${API_KEY_ENV_VAR}`;
const API_KEYS_URL =
  "https://app.umans.ai/billing?context=personal&tab=api-keys";
const DEFAULT_CONTEXT_WINDOW = 128000;
const DEFAULT_MAX_OUTPUT_TOKENS = 32768;
// A Umans API key is static, so the credentials never expire. Use the maximum
// representable Date value so pi never tries to refresh them.
const NEVER_EXPIRES = 8640000000000000;

interface UmansReasoning {
  supported?: boolean;
  can_disable?: boolean;
  levels?: string[];
  default_level?: string | null;
}

interface UmansCapabilities {
  max_completion_tokens?: number;
  recommended_max_tokens?: number;
  context_window?: number;
  supports_vision?: boolean | "via-handoff";
  supports_tools?: boolean;
  reasoning?: UmansReasoning;
}

interface UmansModelInfo {
  name?: string;
  display_name?: string;
  description?: string;
  deprecation?: unknown;
  capabilities?: UmansCapabilities;
}

type UmansModelsInfoResponse = Record<string, UmansModelInfo>;

/**
 * Resolve an output budget that never trips the gateway's hard cap. Umans
 * rejects `max_tokens >= max_completion_tokens` with a 400, so clamp the
 * recommended value to `cap - 1`. Falls back to a sane default when the catalog
 * omits a recommendation.
 */
export function safeMaxTokens(recommended?: number, cap?: number): number {
  let value =
    typeof recommended === "number" && recommended > 0
      ? recommended
      : DEFAULT_MAX_OUTPUT_TOKENS;
  if (typeof cap === "number" && cap > 0) value = Math.min(value, cap - 1);
  return Math.max(value, 1);
}

/**
 * Map pi thinking levels to Umans reasoning levels.
 *
 * Umans exposes levels none/minimal/low/medium/high/xhigh/max; pi exposes
 * off/minimal/low/medium/high/xhigh. pi has no "max", so pi's `xhigh` maps to
 * Umans's `max` when available, giving the deepest tier via pi's highest level.
 * When a model can't disable reasoning (`can_disable === false`), mark pi's
 * `off` as unsupported (null) so pi clamps to the minimum level instead of
 * sending a disabled-thinking parameter the model rejects.
 */
export function toThinkingLevelMap(
  reasoning?: UmansReasoning,
): ProviderModelConfig["thinkingLevelMap"] {
  if (!reasoning?.supported) return {};
  const levels = new Set(reasoning.levels ?? []);
  return {
    off: reasoning.can_disable && levels.has("none") ? "none" : null,
    minimal: levels.has("minimal") ? "minimal" : null,
    low: levels.has("low") ? "low" : null,
    medium: levels.has("medium") ? "medium" : null,
    high: levels.has("high") ? "high" : null,
    xhigh: levels.has("max") ? "max" : levels.has("xhigh") ? "xhigh" : null,
  };
}

/**
 * Turn a catalog entry into a registered model, or `undefined` to skip it.
 * Skips models that don't support tools or that are deprecated. Both native
 * (`supports_vision: true`) and gateway-handoff (`"via-handoff"`) vision models
 * accept images: the Anthropic Messages endpoint handles the handoff
 * server-side, so from pi's side they are image-capable.
 */
export function toRegisteredModel(
  id: string,
  model: UmansModelInfo,
): ProviderModelConfig | undefined {
  const caps = model.capabilities;
  if (!caps?.supports_tools) return undefined;
  if (model.deprecation) return undefined;

  const reasoning = caps.reasoning?.supported ?? false;
  const input: ("text" | "image")[] =
    caps.supports_vision === true || caps.supports_vision === "via-handoff"
      ? ["text", "image"]
      : ["text"];

  return {
    id,
    name: model.display_name ?? model.name ?? id,
    reasoning,
    thinkingLevelMap: toThinkingLevelMap(caps.reasoning),
    input,
    // Umans is a flat subscription; the catalog reports no per-token pricing.
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: caps.context_window ?? DEFAULT_CONTEXT_WINDOW,
    maxTokens: safeMaxTokens(
      caps.recommended_max_tokens,
      caps.max_completion_tokens,
    ),
    compat: {
      // Umans exposes reasoning as effort levels, i.e. Anthropic's adaptive
      // thinking format (`thinking.type: "adaptive"` + `output_config.effort`).
      forceAdaptiveThinking: reasoning,
      // Adaptive thinking returns thinking blocks with no valid signature.
      // pi's default downgrades unsigned prior thinking to text, which stacks a
      // `[Thinking from previous turn]` marker and corrupts context. Replay the
      // block with an empty signature instead — Umans accepts that.
      allowEmptySignature: true,
    },
  };
}

async function fetchModels(): Promise<ProviderModelConfig[] | undefined> {
  const apiKey = process.env[API_KEY_ENV_VAR];
  const headers = apiKey ? { Authorization: `Bearer ${apiKey}` } : undefined;

  try {
    const res = await fetch(MODELS_INFO_URL, { headers });
    if (!res.ok) {
      console.warn(
        `[${PROVIDER_NAME}] API returned ${res.status}: ${res.statusText}`,
      );
      return undefined;
    }

    const response = (await res.json()) as UmansModelsInfoResponse;
    if (typeof response !== "object" || response === null) {
      console.warn(`[${PROVIDER_NAME}] Unexpected API response shape`);
      return undefined;
    }

    return Object.entries(response).flatMap(([id, model]) => {
      const registeredModel = toRegisteredModel(id, model);
      return registeredModel ? [registeredModel] : [];
    });
  } catch (error) {
    console.warn(`[${PROVIDER_NAME}] Failed to fetch models:`, error);
    return undefined;
  }
}

async function login(
  callbacks: OAuthLoginCallbacks,
): Promise<OAuthCredentials> {
  callbacks.onAuth({
    url: API_KEYS_URL,
    instructions:
      "Open your Umans billing page, create an API key, then paste it below.",
  });

  const apiKey = (
    await callbacks.onPrompt({ message: "Paste your Umans API key:" })
  ).trim();
  if (!apiKey) throw new Error("No API key provided");

  return { refresh: apiKey, access: apiKey, expires: NEVER_EXPIRES };
}

export default async function (pi: ExtensionAPI) {
  const models = await fetchModels();
  if (!models) return;

  pi.registerProvider(PROVIDER_NAME, {
    name: PROVIDER_DISPLAY_NAME,
    baseUrl: BASE_URL,
    apiKey: API_KEY_ENV_REF,
    api: "anthropic-messages",
    authHeader: true,
    models,
    oauth: {
      name: PROVIDER_DISPLAY_NAME,
      login,
      refreshToken: async (credentials) => credentials,
      getApiKey: (credentials) => credentials.access,
    },
  });

  pi.registerCommand("umans-models", {
    description: "List available Umans models",
    handler: async (_args, ctx) => {
      if (models.length === 0) {
        ctx.ui.notify("No Umans models available", "warning");
        return;
      }

      const items = [...models]
        .sort((a, b) => a.id.localeCompare(b.id))
        .map((model) => {
          const tags = [];
          if (model.reasoning) tags.push("reasoning");
          if (model.input.includes("image")) tags.push("vision");
          return tags.length > 0
            ? `${model.id} (${tags.join(", ")})`
            : model.id;
        });

      await ctx.ui.select(
        `${PROVIDER_DISPLAY_NAME} — ${models.length} models`,
        items,
      );
    },
  });
}