import { defineExtension } from "@berry/api";
import { type Static, Type } from "@earendil-works/pi-ai";

const parameters = Type.Object({
  sort: Type.Union([
    Type.Literal("default"),
    Type.Literal("throughput"),
    Type.Literal("latency"),
    Type.Literal("price"),
  ], {
    title: "Sort providers",
    description: "throughput is Nitro routing; default preserves the model's routing.",
    default: "default",
  }),
  only: Type.String({
    title: "Allowed providers",
    description:
      "Comma-separated endpoint tags, e.g. baseten/fast. Blank preserves model defaults.",
    default: "",
  }),
  ignore: Type.String({
    title: "Excluded providers",
    description: "Comma-separated provider or endpoint tags.",
    default: "",
  }),
  maxPromptPrice: Type.String({
    title: "Maximum input price ($/M tokens)",
    description: "Nonnegative USD per million tokens; blank preserves model defaults.",
    default: "",
  }),
  maxCompletionPrice: Type.String({
    title: "Maximum output price ($/M tokens)",
    description: "Nonnegative USD per million tokens; blank preserves model defaults.",
    default: "",
  }),
  allowFallbacks: Type.Union([
    Type.Literal("default"),
    Type.Literal("yes"),
    Type.Literal("no"),
  ], {
    title: "Allow provider fallbacks",
    description: "Fallbacks still have to satisfy the allowlist, exclusions, and price caps.",
    default: "default",
  }),
}, { additionalProperties: false });

type State = Static<typeof parameters>;
type Event = { type: "configured"; values: State };
type Routing = {
  sort?: "throughput" | "latency" | "price";
  only?: string[];
  ignore?: string[];
  max_price?: { prompt?: number; completion?: number };
  allow_fallbacks?: boolean;
};

/** The editable fields, since an event recorded before one was dropped still carries it. */
function settings(state: State): State {
  return {
    sort: state.sort,
    only: state.only,
    ignore: state.ignore,
    maxPromptPrice: state.maxPromptPrice,
    maxCompletionPrice: state.maxCompletionPrice,
    allowFallbacks: state.allowFallbacks,
  };
}

function routing(values: State): Routing {
  const result: Routing = {};
  if (values.sort !== "default") result.sort = values.sort;
  for (const key of ["only", "ignore"] as const) {
    const tags = [
      ...new Set(
        values[key].split(",").map((tag) => tag.trim()).filter(Boolean),
      ),
    ];
    if (tags.length) result[key] = tags;
  }
  for (
    const [field, key] of [
      ["maxPromptPrice", "prompt"],
      ["maxCompletionPrice", "completion"],
    ] as const
  ) {
    const text = values[field].trim();
    if (!text) continue;
    const price = Number(text);
    if (!/^(?:\d+(?:\.\d*)?|\.\d+)$/.test(text) || !Number.isFinite(price)) {
      throw new Error(
        `${
          key === "prompt" ? "Input" : "Output"
        } price must be a nonnegative decimal in $/M tokens.`,
      );
    }
    (result.max_price ??= {})[key] = price;
  }
  if (values.allowFallbacks !== "default") {
    result.allow_fallbacks = values.allowFallbacks === "yes";
  }
  return result;
}

function isObject(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

export default defineExtension<State, Event>({
  name: "openrouter-routing",
  description: "per-conversation OpenRouter provider routing and price limits",
  alwaysOn: true,
  sortOrder: -1000,
  initial: () => ({
    sort: "default",
    only: "",
    ignore: "",
    maxPromptPrice: "",
    maxCompletionPrice: "",
    allowFallbacks: "default",
  }),
  apply: (state, event) =>
    event.type === "configured" ? settings({ ...state, ...event.values }) : state,
  setup(ext) {
    ext.configuration({
      parameters,
      values: (state) => settings(state),
      events(values, state) {
        routing(values);
        return Object.keys(values).every((key) =>
            values[key as keyof State] === state[key as keyof State]
          )
          ? []
          : [{ type: "configured", values }];
      },
    });
    return (conv) => {
      const requirements = routing(settings(conv.state));
      if (!Object.keys(requirements).length) return;
      conv.onPayload((payload, model) => {
        if (model.provider !== "openrouter" || !isObject(payload)) return;
        const existing = isObject(payload.provider) ? payload.provider : {};
        const provider = { ...existing, ...requirements };
        if (requirements.max_price) {
          provider.max_price = {
            ...(isObject(existing.max_price) ? existing.max_price : {}),
            ...requirements.max_price,
          };
        }
        return { ...payload, provider };
      });
    };
  },
});