import { StringEnum } from "@earendil-works/pi-ai";
import { type Static, Type } from "typebox";
import type { SearchCommands } from "./types.ts";

const SearchQuerySchema = Type.Object(
  {
    q: Type.String({ description: "Search query." }),
    recency: Type.Optional(
      Type.Integer({ minimum: 0, description: "Only include results from this many recent days." }),
    ),
    domains: Type.Optional(
      Type.Array(Type.String(), { description: "Restrict this query to these domains." }),
    ),
  },
  { additionalProperties: false },
);

const OpenOperationSchema = Type.Object(
  {
    ref_id: Type.String({ description: "Reference ID from an earlier result, or an absolute URL." }),
    lineno: Type.Optional(Type.Integer({ minimum: 0, description: "Line number to position the page at." })),
  },
  { additionalProperties: false },
);

const ClickOperationSchema = Type.Object(
  {
    ref_id: Type.String({ description: "Reference ID containing the numbered link." }),
    id: Type.Integer({ minimum: 0, description: "Numbered link ID to open." }),
  },
  { additionalProperties: false },
);

const FindOperationSchema = Type.Object(
  {
    ref_id: Type.String({ description: "Reference ID or URL to search within." }),
    pattern: Type.String({ description: "Text pattern to find in the page." }),
  },
  { additionalProperties: false },
);

const ScreenshotOperationSchema = Type.Object(
  {
    ref_id: Type.String({ description: "PDF reference ID or URL." }),
    pageno: Type.Integer({ minimum: 0, description: "Zero-indexed PDF page number." }),
  },
  { additionalProperties: false },
);

const FinanceOperationSchema = Type.Object(
  {
    ticker: Type.String({ description: "Ticker symbol." }),
    type: StringEnum(["equity", "fund", "crypto", "index"] as const),
    market: Type.Optional(Type.String({ description: "ISO alpha-3 market, OTC, or empty for crypto." })),
  },
  { additionalProperties: false },
);

const WeatherOperationSchema = Type.Object(
  {
    location: Type.String({ description: "Location in Country, Area, City format." }),
    start: Type.Optional(Type.String({ description: "Start date in YYYY-MM-DD format." })),
    duration: Type.Optional(Type.Integer({ minimum: 1, description: "Number of forecast days." })),
  },
  { additionalProperties: false },
);

const SportsOperationSchema = Type.Object(
  {
    tool: Type.Optional(StringEnum(["sports"] as const)),
    fn: StringEnum(["schedule", "standings"] as const),
    league: StringEnum(["nba", "wnba", "nfl", "nhl", "mlb", "epl", "ncaamb", "ncaawb", "ipl"] as const),
    team: Type.Optional(Type.String()),
    opponent: Type.Optional(Type.String()),
    date_from: Type.Optional(Type.String({ description: "Start date in YYYY-MM-DD format." })),
    date_to: Type.Optional(Type.String({ description: "End date in YYYY-MM-DD format." })),
    num_games: Type.Optional(Type.Integer({ minimum: 1 })),
    locale: Type.Optional(Type.String()),
  },
  { additionalProperties: false },
);

const TimeOperationSchema = Type.Object(
  {
    utc_offset: Type.String({
      pattern: "^[+-](?:0\\d|1\\d|2[0-3]):[0-5]\\d$",
      description: "UTC offset, for example +03:00.",
    }),
  },
  { additionalProperties: false },
);

/** Codex-compatible standalone web-search command schema. */
export const WebSearchParameters = Type.Object(
  {
    search_query: Type.Optional(
      Type.Array(SearchQuerySchema, {
        maxItems: 4,
        description: "Search the internet. Batch up to four independent queries in one call.",
      }),
    ),
    image_query: Type.Optional(Type.Array(SearchQuerySchema, { maxItems: 4, description: "Search for images." })),
    open: Type.Optional(Type.Array(OpenOperationSchema, { description: "Open result references or URLs." })),
    click: Type.Optional(Type.Array(ClickOperationSchema, { description: "Open numbered links from fetched pages." })),
    find: Type.Optional(Type.Array(FindOperationSchema, { description: "Find text within fetched pages." })),
    screenshot: Type.Optional(Type.Array(ScreenshotOperationSchema, { description: "Render PDF pages." })),
    finance: Type.Optional(Type.Array(FinanceOperationSchema, { description: "Look up market prices." })),
    weather: Type.Optional(Type.Array(WeatherOperationSchema, { description: "Look up weather forecasts." })),
    sports: Type.Optional(Type.Array(SportsOperationSchema, { description: "Look up sports schedules or standings." })),
    time: Type.Optional(Type.Array(TimeOperationSchema, { description: "Look up current time by UTC offset." })),
    response_length: Type.Optional(
      StringEnum(["short", "medium", "long"] as const, { description: "Amount of raw search output to return." }),
    ),
  },
  { additionalProperties: false },
);

export type WebSearchInput = Static<typeof WebSearchParameters>;

const ARRAY_KEYS = [
  "search_query",
  "image_query",
  "open",
  "click",
  "find",
  "screenshot",
  "finance",
  "weather",
  "sports",
  "time",
] as const;

export function normalizeSearchCommands(input: WebSearchInput): SearchCommands {
  const normalized: SearchCommands = {};
  for (const key of ARRAY_KEYS) {
    const value = input[key];
    if (value && value.length > 0) Object.assign(normalized, { [key]: value });
  }
  if (input.response_length) normalized.response_length = input.response_length;

  const queryCount = normalized.search_query?.length ?? 0;
  if (queryCount > 4) throw new Error("search_query accepts at most 4 queries per call");
  if (queryCount > 3 && normalized.response_length !== "medium" && normalized.response_length !== "long") {
    throw new Error("response_length must be medium or long when search_query contains 4 queries");
  }
  if (Object.keys(normalized).length === 0) throw new Error("web_search requires at least one search operation");
  return normalized;
}

/** Converts v0.2 `{ query }` calls when a session created by that version is resumed. */
export function prepareWebSearchArguments(args: unknown): unknown {
  if (!args || typeof args !== "object") return args;
  const value = args as { query?: unknown };
  if (typeof value.query !== "string" || !value.query.trim()) return args;
  return {
    search_query: [{ q: value.query.trim() }],
    response_length: "long",
  } satisfies WebSearchInput;
}