import { GoogleGenAI, Type } from "@google/genai";
import { CORRECT_CITY_SCHEMA, CORRECT_CITY_SYSTEM, buildUserMessage, usageCostUsd } from "./prompt.js";

export const name = "gemini";

// Per-million-token rates. Verified against
// https://ai.google.dev/gemini-api/docs/pricing
export const pricingByModel = {
  "gemini-2.0-flash":      { inputUsdPerM: 0.10,  outputUsdPerM: 0.40,  cacheWriteUsdPerM: 0, cacheReadUsdPerM: 0.025 },
  "gemini-2.0-flash-lite": { inputUsdPerM: 0.075, outputUsdPerM: 0.30,  cacheWriteUsdPerM: 0, cacheReadUsdPerM: 0     },
  "gemini-2.5-flash":      { inputUsdPerM: 0.30,  outputUsdPerM: 2.50,  cacheWriteUsdPerM: 0, cacheReadUsdPerM: 0.075 },
};

export const defaultModel = "gemini-2.0-flash";

export function activeModel() {
  return process.env.GEMINI_MODEL || defaultModel;
}

export function isAvailable() {
  return Boolean(process.env.GEMINI_API_KEY || process.env.GOOGLE_API_KEY);
}

let _client = null;
function getClient() {
  if (_client) return _client;
  const apiKey = process.env.GEMINI_API_KEY || process.env.GOOGLE_API_KEY;
  _client = new GoogleGenAI({ apiKey });
  return _client;
}

// Gemini's function declarations want a slightly simplified JSON Schema
// (no `["string", "null"]` union types). Translate at module load.
function toGeminiSchema(node) {
  if (!node || typeof node !== "object") return node;
  const out = { ...node };
  if (Array.isArray(out.type)) {
    out.type = out.type.find((t) => t !== "null") || "string";
    out.nullable = node.type.includes("null");
  }
  if (out.properties) {
    out.properties = Object.fromEntries(
      Object.entries(out.properties).map(([k, v]) => [k, toGeminiSchema(v)]),
    );
  }
  if (out.items) out.items = toGeminiSchema(out.items);
  return out;
}

const TOOL_DECL = {
  name: "correct_city",
  description: "Pick the closest matching US city/town/CDP name from the provided list for the user's input, or return null if nothing reasonable matches.",
  parameters: toGeminiSchema(CORRECT_CITY_SCHEMA),
};

export async function correctCity({ city, stateCode, candidates }) {
  const model = activeModel();
  const pricing = pricingByModel[model] || { inputUsdPerM: 0, outputUsdPerM: 0, cacheWriteUsdPerM: 0, cacheReadUsdPerM: 0 };
  const resp = await getClient().models.generateContent({
    model,
    contents: [{ role: "user", parts: [{ text: buildUserMessage({ city, stateCode, candidates }) }] }],
    config: {
      systemInstruction: CORRECT_CITY_SYSTEM,
      tools: [{ functionDeclarations: [TOOL_DECL] }],
      toolConfig: { functionCallingConfig: { mode: "ANY", allowedFunctionNames: [TOOL_DECL.name] } },
      maxOutputTokens: 200,
    },
  });
  const call = resp.functionCalls?.[0];
  if (!call?.args) throw new Error("Gemini did not return a function call");
  const input = call.args;
  const u = resp.usageMetadata || {};
  const cached = Number(u.cachedContentTokenCount ?? 0);
  const prompt = Number(u.promptTokenCount ?? 0);
  const usage = {
    input_tokens:                Math.max(0, prompt - cached),
    output_tokens:               Number(u.candidatesTokenCount ?? 0),
    cache_creation_input_tokens: 0,
    cache_read_input_tokens:     cached,
  };
  return {
    provider: name,
    model,
    correctedCityName: input.correctedCityName ?? null,
    confidence: input.confidence ?? "low",
    reasoning: input.reasoning ?? "",
    usage,
    costUsd: usageCostUsd(usage, pricing),
    pricingPerMTok: {
      input: pricing.inputUsdPerM,
      output: pricing.outputUsdPerM,
      cacheWrite: pricing.cacheWriteUsdPerM,
      cacheRead: pricing.cacheReadUsdPerM,
    },
  };
}
