import { GoogleGenAI } from "@google/genai";
import fs from "fs";
import { logTrainingTelemetryDb } from "./telemetry-db";

const GEMINI_API_KEY = process.env.GOOGLE_API_KEY || "";
const GEMINI_MODEL = process.env.VIBN_CHAT_MODEL || "gemini-3.1-pro-preview";

// Add a clear visual log so we always know exactly which model is active in the terminal
console.log(`[GeminiChat] Initialized — Model: ${GEMINI_MODEL}`);

if (!GEMINI_API_KEY) {
  console.warn(
    `[GeminiChat] WARNING: GOOGLE_API_KEY is not set. Chat stream will fail with 403 Forbidden.`,
  );
}

const ai = new GoogleGenAI({ apiKey: GEMINI_API_KEY });

export interface ChatMessage {
  role: "user" | "assistant" | "tool";
  content: string;
  toolCalls?: ToolCall[];
  toolCallId?: string;
  toolName?: string;
  thoughtSignature?: string;
}

export interface ToolCall {
  id: string;
  name: string;
  args: Record<string, unknown>;
  thoughtSignature?: string;
}

export interface ToolDefinition {
  name: string;
  description: string;
  parameters: Record<string, unknown>;
}

export interface ChatChunk {
  type: "text" | "thinking" | "tool_call" | "done" | "error";
  text?: string;
  toolCall?: ToolCall;
  error?: string;
}

type GeminiPart = Record<string, unknown>;
type GeminiContent = {
  role: string;
  parts: GeminiPart[];
};

function toGeminiContents(messages: ChatMessage[]) {
  const contents: GeminiContent[] = [];

  for (const msg of messages) {
    if (msg.role === "user") {
      contents.push({ role: "user", parts: [{ text: msg.content }] });
    } else if (msg.role === "assistant") {
      const parts: GeminiPart[] = [];
      if (msg.content) parts.push({ text: msg.content });
      if (msg.toolCalls?.length) {
        for (const tc of msg.toolCalls) {
          const part: GeminiPart = {
            functionCall: { name: tc.name, args: tc.args },
          };
          if (tc.thoughtSignature) {
            part.thoughtSignature = tc.thoughtSignature;
          }
          parts.push(part);
        }
      }
      if (parts.length) contents.push({ role: "model", parts });
    } else if (msg.role === "tool") {
      const part = {
        functionResponse: {
          name: msg.toolName || "unknown",
          response: { name: msg.toolName || "unknown", content: msg.content },
        },
      };
      const last = contents[contents.length - 1];
      if (last?.role === "user") {
        last.parts.push(part);
      } else {
        contents.push({ role: "user", parts: [part] });
      }
    }
  }
  return contents;
}

function toGeminiFunctions(tools: ToolDefinition[]) {
  if (!tools.length) return undefined;
  return [
    {
      functionDeclarations: tools.map((t) => ({
        name: t.name,
        description: t.description,
        parameters: t.parameters as Record<string, unknown>,
      })),
    },
  ];
}

export async function callGeminiChat(opts: {
  systemPrompt: string;
  messages: ChatMessage[];
  tools?: ToolDefinition[];
  temperature?: number;
  includeThoughts?: boolean;
  /** Cancels the in-flight generation when the user clicks Stop. */
  signal?: AbortSignal;
}): Promise<{
  text: string;
  thoughts: string;
  toolCalls: ToolCall[];
  finishReason?: string;
  error?: string;
}> {
  try {
    const config: Record<string, unknown> = {
      temperature: opts.temperature ?? 0.7,
      maxOutputTokens: 8192,
    };

    if (opts.systemPrompt) {
      config.systemInstruction = opts.systemPrompt;
    }

    const fns = toGeminiFunctions(opts.tools ?? []);
    if (fns) config.tools = fns;

    // Forward the request abort signal so a Stop click cancels the actual
    // Gemini generation instead of only short-circuiting the next round.
    if (opts.signal) config.abortSignal = opts.signal;

    console.log("\n========================================================");
    console.log("➡️ [GEMINI API REQUEST]");
    console.log("========================================================");
    console.log(
      `System Prompt: ${typeof config.systemInstruction === "string" ? config.systemInstruction.slice(0, 1000) + "..." : "None"}`,
    );
    console.log(
      "Contents Payload:",
      JSON.stringify(toGeminiContents(opts.messages), null, 2),
    );
    console.log("========================================================\n");

    const startTime = Date.now();
    const response = await ai.models.generateContent({
      model: GEMINI_MODEL,
      contents: toGeminiContents(opts.messages),
      config,
    });
    const durationMs = Date.now() - startTime;

    console.log("\n========================================================");
    console.log("⬅️ [GEMINI API RESPONSE]");
    console.log("========================================================");
    console.log(
      "Raw Candidates:",
      JSON.stringify(response.candidates, null, 2),
    );
    console.log("========================================================\n");

    console.log(
      "[GeminiChat] Raw Response:",
      JSON.stringify(response, null, 2),
    );

    try {
      fs.writeFileSync(
        "/tmp/last_gemini.json",
        JSON.stringify(response, null, 2),
        "utf8",
      );
    } catch (e) {
      console.warn("Failed to write debug response to disk:", e);
    }

    let text = "";
    let thoughts = "";
    const toolCalls: ToolCall[] = [];

    const parts = response.candidates?.[0]?.content?.parts ?? [];
    const isPartThought = (p: Record<string, unknown>) =>
      Boolean(p.thought || p.thoughtSignature);

    for (const part of parts) {
      if (part.text) {
        if (isPartThought(part as Record<string, unknown>))
          thoughts += part.text;
        else text += part.text;
      }
      if (part.functionCall) {
        toolCalls.push({
          id: `tc-${Date.now()}-${Math.random().toString(36).slice(2)}`,
          name: part.functionCall.name,
          args: (part.functionCall.args as Record<string, unknown>) ?? {},
          thoughtSignature: (part as { thoughtSignature?: string })
            .thoughtSignature,
        });
      }
    }

    // Empty-answer fallback: Gemini "thinking" responses can spend the whole
    // turn emitting `thought` parts and return NO answer part — leaving `text`
    // empty. Mirror the OpenAI-compatible adapter (which promotes
    // reasoning_content when content is empty) so the user never gets a blank
    // bubble. Only promote when there's also no tool call to render.
    if (!text.trim() && thoughts.trim() && toolCalls.length === 0) {
      text = thoughts.trim();
    }

    logTrainingTelemetryDb({
      model: GEMINI_MODEL,
      systemPrompt: opts.systemPrompt,
      messages: opts.messages,
      response: { text, thoughts, toolCalls },
      metrics: {
        promptTokens: response.usageMetadata?.promptTokenCount,
        completionTokens: response.usageMetadata?.candidatesTokenCount,
        totalTokens: response.usageMetadata?.totalTokenCount,
        durationMs,
      },
    });

    return {
      text,
      thoughts,
      toolCalls,
      finishReason: response.candidates?.[0]?.finishReason,
    };
  } catch (error) {
    const aborted =
      opts.signal?.aborted ||
      (error instanceof Error && error.name === "AbortError");
    return {
      text: "",
      thoughts: "",
      toolCalls: [],
      error: aborted
        ? "aborted"
        : `GoogleGenAI error: ${error instanceof Error ? error.message : String(error)}`,
    };
  }
}

export async function* streamGeminiChat(opts: {
  systemPrompt: string;
  messages: ChatMessage[];
  tools?: ToolDefinition[];
  temperature?: number;
}): AsyncGenerator<ChatChunk> {
  try {
    const config: Record<string, unknown> = {
      temperature: opts.temperature ?? 0.7,
      maxOutputTokens: 8192,
    };

    if (opts.systemPrompt) {
      config.systemInstruction = opts.systemPrompt;
    }

    const fns = toGeminiFunctions(opts.tools ?? []);
    if (fns) config.tools = fns;

    console.log("\n========================================================");
    console.log("➡️ [GEMINI STREAM REQUEST]");
    console.log("========================================================");
    console.log(
      `System Prompt: ${typeof config.systemInstruction === "string" ? config.systemInstruction.slice(0, 1000) + "..." : "None"}`,
    );
    console.log(
      "Contents Payload:",
      JSON.stringify(toGeminiContents(opts.messages), null, 2),
    );
    console.log("========================================================\n");

    const startTime = Date.now();
    const streamResult = await ai.models.generateContentStream({
      model: GEMINI_MODEL,
      contents: toGeminiContents(opts.messages),
      config,
    });

    console.log("[GeminiChat] Stream request initiated");
    const isPartThought = (p: Record<string, unknown>) =>
      Boolean(p.thought || p.thoughtSignature);

    let text = "";
    let thoughts = "";
    const toolCalls: any[] = [];

    for await (const chunk of streamResult) {
      const parts = chunk.candidates?.[0]?.content?.parts ?? [];
      for (const part of parts) {
        if (part.text) {
          if (isPartThought(part as Record<string, unknown>)) {
            thoughts += part.text;
            yield { type: "thinking", text: part.text };
          } else {
            text += part.text;
            yield { type: "text", text: part.text };
          }
        }
        if (part.functionCall) {
          toolCalls.push(part.functionCall);
        }
      }
    }
    const durationMs = Date.now() - startTime;

    logTrainingTelemetryDb({
      model: GEMINI_MODEL,
      systemPrompt: opts.systemPrompt,
      messages: opts.messages,
      response: { text, thoughts, toolCalls },
      metrics: {
        durationMs,
      },
    });

    yield { type: "done" };
  } catch (error) {
    yield {
      type: "error",
      error: `GoogleGenAI error: ${error instanceof Error ? error.message : String(error)}`,
    };
  }
}
