import OpenAI from "openai";
import { Client } from "@gradio/client";
import { LIGHTNING_BASE } from "./config.js";

// ── Gradio search client singleton ─────────────────────────────────────────
// Connect once at module load and reuse across all searches.
// If the connection drops, gradioClient is reset to null so the next call
// reconnects automatically rather than retrying a dead client forever.
let gradioClient = null;
let gradioConnecting = null; // in-flight connect promise, prevents thundering herd

async function getGradioClient() {
  if (gradioClient) return gradioClient;
  if (gradioConnecting) return gradioConnecting;
  gradioConnecting = Client.connect("incognitolm/Web-Search")
    .then(c  => { gradioClient = c; gradioConnecting = null; return c; })
    .catch(e => { gradioConnecting = null; throw e; });
  return gradioConnecting;
}

async function gradioSearch(query) {
  // Try once; if it fails reset the singleton so the next call gets a fresh connection.
  try {
    const client = await getGradioClient();
    const r = await client.predict("/perform_search", { query });
    // r.data is an array; the search results are in the first element.
    // Accept a string, an array, or an object — normalise all to a string.
    const raw = Array.isArray(r.data) ? r.data[0] : r.data;
    if (!raw) throw new Error("Empty response from search endpoint");
    return typeof raw === "string" ? raw : JSON.stringify(raw);
  } catch (err) {
    // Invalidate the singleton so the next search attempt reconnects.
    gradioClient = null;
    throw err;
  }
}

const SYSTEM_PROMPT =
  "CRITICAL RULE: Every response MUST use HTML <span data-color=\"{COLOR NAME}\"> tags to color main points and headings. " +
  "COLORS MUST HAVE MEANING AND CONSISTENCY ACROSS THE ENTIRE CONVERSATION. " +
  "You may ONLY use the following semantic color names: green, pink, blue, red, orange, yellow, purple, teal, gold, coral. " +
  "Never output text formatted with explicit black or white colors. " +
  "Use a variety of colors throughout every response to distinguish headings, sections, and key terms. " +
  "Keep code blocks plain, but color headings and important points in surrounding text. " +
  "Do not over-color responses. Use color intentionally and sparingly. " +
  "CRITICAL RULE: MARKDOWN FORMATTING SUCH AS #, ##, ###, **, * MUST BE PLACED OUTSIDE <span> tags. " +
  "You are a helpful, friendly AI assistant. Use tools when appropriate to help the user. " +
  "When generating media, do not include URLs — it is displayed automatically. " +
  "You can render SVG images by outputting SVG code in a code block tagged exactly as:\n```svg\n<svg>...</svg>\n```\n" +
  "Never use single backslashes. You may use emojis where appropriate. " +
  "Use markdown for everything other than coloring your text. Use tables, lists, and other markdown elements.";

/**
 * Build a per-request OpenAI client pointed at the Lightning backend.
 * A new client is created each call so per-user auth headers are always fresh.
 */
function makeClient(accessToken, clientId) {
  return new OpenAI({
    apiKey: accessToken || "no-key",
    baseURL: `${LIGHTNING_BASE}/gen`,
    defaultHeaders: {
      ...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
      ...(clientId    ? { "X-Client-ID": clientId }                : {}),
    },
  });
}

/**
 * Consume an OpenAI streaming response, firing onToken for each text delta
 * and collecting any tool-call chunks into a finished toolCalls array.
 * Returns { assistantText, toolCalls }.
 */
async function consumeStream(stream, onToken) {
  let assistantText = "";
  const toolCallBuffer = new Map();

  for await (const chunk of stream) {
    const delta = chunk.choices?.[0]?.delta;
    if (!delta) continue;

    if (delta.content) {
      assistantText += delta.content;
      onToken(delta.content);
    }

    if (delta.tool_calls) {
      for (const call of delta.tool_calls) {
        const entry = toolCallBuffer.get(call.index) ?? { arguments: "" };
        if (call.id)               entry.id   = call.id;
        if (call.function?.name)   entry.name = call.function.name;
        if (call.function?.arguments) entry.arguments += call.function.arguments;
        toolCallBuffer.set(call.index, entry);
      }
    }
  }

  const toolCalls = [...toolCallBuffer.values()].map(t => ({
    id:       t.id || `call_${crypto.randomUUID()}`,
    type:     "function",
    function: { name: t.name, arguments: t.arguments },
  }));

  return { assistantText, toolCalls };
}

export async function streamChat(ws, {
  sessionId,
  model,
  history,
  userMessage,
  tools,
  accessToken,
  clientId,
  onToken,
  onDone,
  onError,
  onToolCall,
  onNewAsset,
  abortSignal,
}) {
  const client       = makeClient(accessToken, clientId);
  const enabledTools = buildToolList(tools);

  const messages = [
    { role: "system", content: SYSTEM_PROMPT },
    ...history.map(normalizeMessage).filter(Boolean),
    { role: "user", content: userMessage },
  ];

  try {
    // ── First stream ────────────────────────────────────────────────────────
    const stream = await client.chat.completions.create({
      model:  model || "lightning",
      messages,
      tools:  enabledTools.length > 0 ? enabledTools : undefined,
      stream: true,
    }, { signal: abortSignal });

    let { assistantText, toolCalls } = await consumeStream(stream, onToken);

    // ── Tool calls → follow-up stream ───────────────────────────────────────
    if (toolCalls.length > 0) {
      const toolResults = await processToolCalls(
        ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset,
      );

      const followUpMessages = [
        { role: "system",    content: SYSTEM_PROMPT },
        ...history.map(normalizeMessage).filter(Boolean),
        { role: "user",      content: userMessage },
        { role: "assistant", content: assistantText || "", tool_calls: toolCalls },
        ...toolResults,
      ];

      const followUpStream = await client.chat.completions.create({
        model:  model || "lightning",
        messages: followUpMessages,
        stream: true,
      }, { signal: abortSignal });

      const followUp = await consumeStream(followUpStream, onToken);
      assistantText += followUp.assistantText;
    }

    onDone(assistantText, toolCalls);
  } catch (err) {
    if (err.name === "AbortError" || err.constructor?.name === "APIUserAbortError") {
      onDone(null, null, true); // aborted
    } else {
      onError(String(err));
    }
  }
}

const VALID_ROLES = new Set(["system", "user", "assistant", "tool"]);

function normalizeMessage(msg) {
  // Drop asset entries (role: "image"/"video"/"audio") — these are UI-only
  // and sending them to the LLM causes invalid-role rejections / blank responses
  if (!VALID_ROLES.has(msg.role)) return null;

  if (msg.role === "assistant" && msg.tool_calls) {
    return { role: "assistant", content: "", tool_calls: msg.tool_calls };
  }
  // Flatten multipart content arrays (e.g. image attachments) to text-only for history
  if (Array.isArray(msg.content)) {
    const textOnly = msg.content
      .filter(b => b.type === "text")
      .map(b => b.text)
      .join("\n");
    return { role: msg.role, content: textOnly || "" };
  }
  return { role: msg.role, content: msg.content };
}

function buildToolList(tools) {
  if (!tools) return [];
  const list = [];
  if (tools.webSearch) {
    list.push({
      type: "function",
      function: {
        name: "ollama_search",
        description: "Search the web for current information",
        parameters: {
          type: "object",
          properties: { query: { type: "string", description: "Search query" } },
          required: ["query"],
        },
      },
    });
    list.push({
      type: "function",
      function: {
        name: "read_web_page",
        description: "Read the content of a web page by URL",
        parameters: {
          type: "object",
          properties: { url: { type: "string", description: "URL to fetch" } },
          required: ["url"],
        },
      },
    });
  }
  if (tools.imageGen) {
    list.push({
      type: "function",
      function: {
        name: "generate_image",
        description: "Generate an image from a prompt",
        parameters: {
          type: "object",
          properties: {
            prompt: { type: "string" },
            mode: { type: "string", enum: ["auto", "fantasy", "realistic"] },
            image_urls: { type: "array", items: { type: "string" } },
          },
          required: ["prompt"],
        },
      },
    });
  }
  if (tools.videoGen) {
    list.push({
      type: "function",
      function: {
        name: "generate_video",
        description: "Generate a video from a prompt",
        parameters: {
          type: "object",
          properties: {
            prompt: { type: "string" },
            ratio: { type: "string", enum: ["3:2", "2:3", "1:1"] },
            mode: { type: "string", enum: ["normal", "fun"] },
            duration: { type: "number" },
            image_urls: { type: "array", items: { type: "string" } },
          },
          required: ["prompt"],
        },
      },
    });
  }
  if (tools.audioGen) {
    list.push({
      type: "function",
      function: {
        name: "generate_audio",
        description: "Generate music or sound effects from a prompt",
        parameters: {
          type: "object",
          properties: { prompt: { type: "string" } },
          required: ["prompt"],
        },
      },
    });
  }
  return list;
}

async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset) {
  const toolResults = [];
  const authHeaders = {};
  if (accessToken) authHeaders["Authorization"] = `Bearer ${accessToken}`;
  if (clientId) authHeaders["X-Client-ID"] = clientId;

  for (const call of toolCalls) {
    let args;
    try { args = JSON.parse(call.function.arguments || "{}"); } catch { args = {}; }

    onToolCall({ id: call.id, name: call.function.name, state: "pending", args });

    let result = "Tool completed.";

    try {
      if (call.function.name === "ollama_search") {
        result = await gradioSearch(args.query);
      }

      else if (call.function.name === "read_web_page") {
        const { convert } = await import("html-to-text");
        const res = await fetch(args.url, { signal: abortSignal });
        if (!res.ok) {
          result = `Failed to fetch: ${res.status}`;
        } else {
          const html = await res.text();
          const titleMatch = html.match(/<title>(.*?)<\/title>/i);
          result = JSON.stringify({
            title: titleMatch?.[1] || "No title",
            content: convert(html, { wordwrap: false }).slice(0, 8000),
          });
        }
      }

      else if (call.function.name === "generate_image") {
        const body = { prompt: args.prompt };
        if (args.mode) body.mode = args.mode;
        if (args.image_urls?.length) body.image_urls = args.image_urls;

        const res = await fetch(`${LIGHTNING_BASE}/gen/image`, {
          method: "POST",
          headers: { "Content-Type": "application/json", ...authHeaders },
          body: JSON.stringify(body),
          signal: abortSignal,
        });
        if (res.ok) {
          const buf = await res.arrayBuffer();
          const ct = res.headers.get("content-type") || "image/png";
          const b64 = Buffer.from(buf).toString("base64");
          const dataUrl = `data:${ct};base64,${b64}`;
          onNewAsset({ role: "image", content: dataUrl });
          result = "Image generated successfully and shown to the user.";
        } else {
          result = `Image generation failed: ${res.status}`;
        }
      }

      else if (call.function.name === "generate_video") {
        const body = { prompt: args.prompt };
        if (args.ratio) body.ratio = args.ratio;
        if (args.mode) body.mode = args.mode;
        if (args.duration) body.duration = args.duration;
        if (args.image_urls?.length) body.image_urls = args.image_urls;

        const res = await fetch(`${LIGHTNING_BASE}/gen/video`, {
          method: "POST",
          headers: { "Content-Type": "application/json", ...authHeaders },
          body: JSON.stringify(body),
          signal: abortSignal,
        });
        if (res.ok) {
          const buf = await res.arrayBuffer();
          const b64 = Buffer.from(buf).toString("base64");
          const dataUrl = `data:video/mp4;base64,${b64}`;
          onNewAsset({ role: "video", content: dataUrl });
          result = "Video generated successfully and shown to the user.";
        } else {
          result = `Video generation failed: ${res.status}`;
        }
      }

      else if (call.function.name === "generate_audio") {
        const res = await fetch(`${LIGHTNING_BASE}/gen/sfx`, {
          method: "POST",
          headers: { "Content-Type": "application/json", ...authHeaders },
          body: JSON.stringify({ prompt: args.prompt }),
          signal: abortSignal,
        });
        if (res.ok) {
          const buf = await res.arrayBuffer();
          const b64 = Buffer.from(buf).toString("base64");
          const dataUrl = `data:audio/mpeg;base64,${b64}`;
          onNewAsset({ role: "audio", content: dataUrl });
          result = "Audio generated successfully and shown to the user.";
        } else {
          result = `Audio generation failed: ${res.status}`;
        }
      }
    } catch (err) {
      console.log(`Tool error: ${String(err)}`);
      result = `Tool error: ${String(err)}`;
    }

    onToolCall({ id: call.id, name: call.function.name, state: "resolved", result });

    toolResults.push({
      role: "tool",
      tool_call_id: call.id,
      content: typeof result === "string" ? result : JSON.stringify(result),
    });
  }

  return toolResults;
}