import express from "express";
import swaggerUi from "swagger-ui-express";
import {
  AutoProcessor,
  Qwen3_5ForConditionalGeneration,
  RawImage,
} from "@huggingface/transformers";
import crypto from "crypto";
import { Readable } from "node:stream";
import { pipeline } from "node:stream/promises";
import { NodeSDK } from "@opentelemetry/sdk-node";
import { LangfuseSpanProcessor } from "@langfuse/otel";
import { startActiveObservation } from "@langfuse/tracing";

const app = express();
const PORT = 7860;
const MODEL_ID = "huggingworld/Qwen3.5-0.8B-ONNX";
const API_KEY = process.env.API_KEY;
const NVIDIA_CHAT_URL =
  process.env.NVIDIA_CHAT_URL || "https://integrate.api.nvidia.com/v1/chat/completions";
const NVIDIA_API_KEY = process.env.NVIDIA_API_KEY;
const NVIDIA_FETCH_TIMEOUT_MS = parseInt(process.env.NVIDIA_FETCH_TIMEOUT_MS || "300000", 10);
const LANGFUSE_PUBLIC_KEY = process.env.LANGFUSE_PUBLIC_KEY;
const LANGFUSE_SECRET_KEY = process.env.LANGFUSE_SECRET_KEY;
const LANGFUSE_BASE_URL = process.env.LANGFUSE_BASE_URL || "https://cloud.langfuse.com";
const LANGFUSE_ENV = process.env.LANGFUSE_ENV || process.env.NODE_ENV || "development";
const LANGFUSE_ENABLED = Boolean(LANGFUSE_PUBLIC_KEY && LANGFUSE_SECRET_KEY);

let model = null;
let processor = null;
let inferenceQueue = Promise.resolve();
let telemetrySdk = null;

function log(level, event, meta = {}) {
  const payload = {
    ts: new Date().toISOString(),
    level,
    event,
    ...meta,
  };
  const line = JSON.stringify(payload);
  if (level === "error") {
    console.error(line);
  } else {
    console.log(line);
  }
}

async function loadModel() {
  const start = Date.now();
  log("info", "model_load_started", { model_id: MODEL_ID });
  processor = await AutoProcessor.from_pretrained(MODEL_ID);
  model = await Qwen3_5ForConditionalGeneration.from_pretrained(MODEL_ID, {
    dtype: {
      embed_tokens: "q4",
      vision_encoder: "fp16",
      decoder_model_merged: "q4",
    },
  });
  log("info", "model_load_completed", {
    model_id: MODEL_ID,
    duration_ms: Date.now() - start,
  });
}

async function setupTracing() {
  if (!LANGFUSE_ENABLED) {
    log("info", "langfuse_disabled", {
      reason: "missing_langfuse_keys",
    });
    return;
  }

  const spanProcessor = new LangfuseSpanProcessor({
    publicKey: LANGFUSE_PUBLIC_KEY,
    secretKey: LANGFUSE_SECRET_KEY,
    baseUrl: LANGFUSE_BASE_URL,
    environment: LANGFUSE_ENV,
  });

  telemetrySdk = new NodeSDK({
    spanProcessors: [spanProcessor],
  });

  await telemetrySdk.start();
  log("info", "langfuse_enabled", {
    base_url: LANGFUSE_BASE_URL,
    environment: LANGFUSE_ENV,
  });
}

async function withPromptTrace(req, prompt, maxTokens, handler) {
  if (!LANGFUSE_ENABLED) {
    return handler();
  }

  return startActiveObservation(
    "http.prompt",
    async (span) => {
      span.update({
        input: { prompt, max_tokens: maxTokens },
        metadata: {
          request_id: req.requestId,
          method: req.method,
          path: req.originalUrl,
        },
      });

      return handler(span);
    },
    { endOnExit: true },
  );
}

function queueInference(runner) {
  const task = inferenceQueue.then(() => runner());
  inferenceQueue = task.catch(() => {});
  return task;
}

async function runTextInference(prompt, maxTokens) {
  const conversation = [
    {
      role: "user",
      content: [{ type: "text", text: prompt }],
    },
  ];

  const text = processor.apply_chat_template(conversation, {
    add_generation_prompt: true,
    chat_template_kwargs: { enable_thinking: false },
  });

  const inputs = await processor(text);
  const output = await model.generate({
    ...inputs,
    max_new_tokens: maxTokens,
    do_sample: false,
  });

  const promptLength = inputs.input_ids.dims.at(-1);
  const decoded = processor.batch_decode(
    output.slice(null, [promptLength, null]),
    { skip_special_tokens: true },
  );
  return decoded[0];
}

function queueTextInference(prompt, maxTokens) {
  return queueInference(() => runTextInference(prompt, maxTokens));
}

function decodeDataUrl(url) {
  const trimmed = url.trim();
  const m = /^data:([^;,]+)?;base64,(.+)$/i.exec(trimmed);
  if (!m) return null;
  const mime = (m[1] || "application/octet-stream").split(";")[0].trim();
  try {
    return { mime, buffer: Buffer.from(m[2], "base64") };
  } catch {
    return null;
  }
}

async function rawImageFromDataUrlOnly(url) {
  const data = decodeDataUrl(url);
  if (!data) {
    return null;
  }
  return RawImage.fromBlob(new Blob([data.buffer], { type: data.mime }));
}

async function runVisionInference(userText, rawImages, maxTokens) {
  const content = [];
  for (let i = 0; i < rawImages.length; i++) {
    content.push({ type: "image" });
  }

  const textLine = (userText ?? "").trim() || "What is in this image?";
  content.push({ type: "text", text: textLine });

  const conversation = [
    {
      role: "user",
      content,
    },
  ];

  const text = processor.apply_chat_template(conversation, {
    add_generation_prompt: true,
    chat_template_kwargs: { enable_thinking: false },
  });

  const visionArg = rawImages.length === 1 ? rawImages[0] : rawImages;
  const inputs = await processor(text, visionArg);
  const output = await model.generate({
    ...inputs,
    max_new_tokens: maxTokens,
    do_sample: false,
  });

  const promptLength = inputs.input_ids.dims.at(-1);
  const decoded = processor.batch_decode(output.slice(null, [promptLength, null]), {
    skip_special_tokens: true,
  });
  return decoded[0];
}

function getLastUserMessage(messages) {
  if (!Array.isArray(messages)) return null;
  for (let i = messages.length - 1; i >= 0; i--) {
    if (messages[i]?.role === "user") return messages[i];
  }
  return null;
}

function flattenUserMultimodalContent(userMessage) {
  if (!userMessage) return { segments: [] };
  const c = userMessage.content;
  /** @type {{ type: string, text?: string, url?: string }[]} */
  const segments = [];

  if (typeof c === "string") {
    segments.push({ type: "text", text: c });
    return { segments };
  }

  if (!Array.isArray(c)) return { segments };

  for (const part of c) {
    if (!part || typeof part !== "object") continue;
    if (part.type === "text" && typeof part.text === "string") {
      segments.push({ type: "text", text: part.text });
    } else if (part.type === "image_url") {
      const imgUrl = part.image_url?.url;
      if (typeof imgUrl === "string" && imgUrl.length > 0) {
        segments.push({ type: "image_url", url: imgUrl });
      }
    }
  }

  return { segments };
}

/** @returns {Promise<boolean>} true if the response was fully sent (do not send fallback) */
async function proxyNvidiaChatCompletions(req, body, streamRequested, res) {
  if (!NVIDIA_API_KEY) {
    log("info", "nvidia_skipped", { request_id: req.requestId, reason: "no_nvidia_api_key" });
    return false;
  }

  const controller = new AbortController();
  const timeoutId = setTimeout(() => controller.abort(), NVIDIA_FETCH_TIMEOUT_MS);
  try {
    const upstream = await fetch(NVIDIA_CHAT_URL, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${NVIDIA_API_KEY}`,
        Accept: streamRequested ? "text/event-stream" : "application/json",
      },
      body: JSON.stringify(body),
      signal: controller.signal,
    });

    if (!upstream.ok) {
      let preview = "";
      try {
        preview = (await upstream.text()).slice(0, 800);
      } catch {
        /* ignore */
      }
      log("warn", "nvidia_upstream_rejected", {
        request_id: req.requestId,
        status: upstream.status,
        body_preview: preview,
      });
      return false;
    }

    const ct = upstream.headers.get("content-type") || "";

    if (upstream.body != null && streamRequested && ct.includes("text/event-stream")) {
      res.status(200);
      res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
      res.setHeader("Cache-Control", "no-cache");
      res.setHeader("Connection", "keep-alive");
      try {
        await pipeline(
          Readable.fromWeb(/** @type {any} */ (upstream.body)),
          res,
        );
      } catch (pipeErr) {
        log("error", "nvidia_stream_pipe_failed", {
          request_id: req.requestId,
          error: pipeErr.message,
        });
        if (!res.headersSent) {
          return false;
        }
        if (!res.writableEnded) {
          try {
            res.destroy();
          } catch {
            /* ignore */
          }
        }
      }
      return true;
    }

    const buf = await upstream.text();
    if (streamRequested && ct.includes("text/event-stream")) {
      res.status(200);
      res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
      res.send(buf);
      return true;
    }
    res.status(200).type(ct || "application/json").send(buf);
    return true;
  } catch (err) {
    log("warn", "nvidia_fetch_failed", {
      request_id: req.requestId,
      error: err.message,
      name: err.name,
    });
    return false;
  } finally {
    clearTimeout(timeoutId);
  }
}

function openAiCompletionJson(assistantText) {
  const id = `chatcmpl-${crypto.randomUUID()}`;
  return {
    id,
    object: "chat.completion",
    created: Math.floor(Date.now() / 1000),
    model: MODEL_ID,
    choices: [
      {
        index: 0,
        message: { role: "assistant", content: assistantText },
        finish_reason: "stop",
      },
    ],
  };
}

function writeFallbackSseChunks(res, assistantText) {
  const id = `chatcmpl-${crypto.randomUUID()}`;
  const base = () => ({
    id,
    object: "chat.completion.chunk",
    created: Math.floor(Date.now() / 1000),
    model: MODEL_ID,
  });

  res.write(`data: ${JSON.stringify({
    ...base(),
    choices: [{ index: 0, delta: { role: "assistant" }, finish_reason: null }],
  })}\n\n`);

  res.write(`data: ${JSON.stringify({
    ...base(),
    choices: [{ index: 0, delta: { content: assistantText }, finish_reason: null }],
  })}\n\n`);

  res.write(`data: ${JSON.stringify({
    ...base(),
    choices: [{ index: 0, delta: {}, finish_reason: "stop" }],
  })}\n\n`);

  res.write("data: [DONE]\n\n");
}

async function fallbackLocalChatCompletion(req, res, body, streamRequested) {
  if (!model || !processor) {
    log("error", "fallback_model_unavailable", { request_id: req.requestId });
    return res.status(503).json({
      error: { message: "Local ONNX model is not loaded." },
      source: "local_onnx_unavailable",
    });
  }

  const rawMax = typeof body.max_tokens === "number" ? body.max_tokens : 512;
  const maxTokens =
    typeof rawMax === "number" && rawMax > 0 ? Math.min(Math.floor(rawMax), 16384) : 512;

  const lastUser = getLastUserMessage(body.messages);
  const { segments } = flattenUserMultimodalContent(lastUser);

  const textParts = segments
    .filter((s) => s.type === "text")
    .map((s) => s.text ?? "")
    .join("\n");
  const imageUrls = segments.filter((s) => s.type === "image_url");

  let assistantText;
  try {
    if (imageUrls.length > 0) {
      const unsupported = imageUrls.find((s) => !decodeDataUrl(s.url));
      if (unsupported) {
        return res.status(400).json({
          error: {
            message:
              "Local ONNX fallback only accepts embedded images as data:image/...;base64,... URIs.",
          },
          source: "local_onnx_fallback",
        });
      }
      const rawImages = await Promise.all(imageUrls.map((s) => rawImageFromDataUrlOnly(s.url)));
      assistantText = await queueInference(() =>
        runVisionInference(textParts || "What is in this image?", rawImages, maxTokens),
      );
    } else if ((textParts ?? "").trim().length > 0) {
      assistantText = await queueInference(() => runTextInference(textParts.trim(), maxTokens));
    } else {
      return res.status(400).json({
        error: {
          message:
            "Local fallback needs text or image parts on the last user message.",
        },
        source: "local_onnx_fallback",
      });
    }

    log("info", "local_fallback_completed", {
      request_id: req.requestId,
      chars: assistantText?.length ?? 0,
      had_images: imageUrls.length > 0,
    });

    if (streamRequested) {
      res.status(200);
      res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
      res.setHeader("Cache-Control", "no-cache");
      res.setHeader("Connection", "keep-alive");
      writeFallbackSseChunks(res, assistantText);
      return res.end();
    }

    return res.json(openAiCompletionJson(assistantText));
  } catch (err) {
    log("error", "local_fallback_failed", {
      request_id: req.requestId,
      error: err.message,
      stack: err.stack,
    });
    return res.status(500).json({
      error: { message: err.message || "Local ONNX inference failed." },
      source: "local_onnx_fallback",
    });
  }
}

const swaggerDoc = {
  openapi: "3.0.0",
  info: {
    title: "Qwen3.5-0.8B Text API (ONNX)",
    version: "1.1.0",
    description:
      "/prompt serves local text. POST /v1/chat/completions tries NVIDIA integrate.api (NVIDIA_API_KEY) then falls back to local ONNX (vision + text) on upstream errors or timeouts.",
  },
  components: {
    securitySchemes: {
      ApiKeyAuth: {
        type: "apiKey",
        in: "header",
        name: "X-API-Key",
        description: "Set API_KEY env var; send as X-API-Key or Authorization: Bearer <key>",
      },
    },
  },
  paths: {
    "/": {
      get: {
        summary: "Root",
        responses: { 200: { description: "API status" } },
      },
    },
    "/health": {
      get: {
        summary: "Health check",
        responses: { 200: { description: "Model load status" } },
      },
    },
    "/prompt": {
      post: {
        summary: "Text prompt inference (no image)",
        requestBody: {
          required: true,
          content: {
            "application/json": {
              schema: {
                type: "object",
                required: ["prompt"],
                properties: {
                  prompt: { type: "string", description: "Text prompt to send to the model" },
                  max_tokens: { type: "integer", default: 256 },
                },
              },
            },
          },
        },
        responses: {
          200: { description: "Inference result" },
          400: { description: "Invalid input" },
          401: { description: "Invalid or missing API key" },
          503: { description: "Model not loaded" },
        },
        security: [{ ApiKeyAuth: [] }],
      },
    },
    "/v1/chat/completions": {
      post: {
        summary: "OpenAI-compatible chat completions; NVIDIA proxy with local ONNX fallback",
        responses: {
          200: { description: "Completion JSON or streamed events" },
          400: { description: "Invalid input" },
          401: { description: "Invalid or missing API key" },
          503: { description: "Model not loaded during fallback" },
        },
        security: [{ ApiKeyAuth: [] }],
      },
    },
  },
};

function requireApiKey(req, res, next) {
  if (!API_KEY) return next();
  const bearer = req.headers.authorization?.startsWith("Bearer ")
    ? req.headers.authorization.slice(7)
    : null;
  const key = bearer ?? req.headers["x-api-key"] ?? null;
  if (key !== API_KEY) {
    log("warn", "api_key_rejected", { request_id: req.requestId, path: req.path });
    return res.status(401).json({ detail: "Invalid or missing API key." });
  }
  next();
}

app.use("/docs", swaggerUi.serve, swaggerUi.setup(swaggerDoc));

app.use((req, res, next) => {
  const requestId = crypto.randomUUID();
  const start = Date.now();
  req.requestId = requestId;
  log("info", "request_started", {
    request_id: requestId,
    method: req.method,
    path: req.originalUrl,
    ip: req.ip,
  });
  res.on("finish", () => {
    log("info", "request_finished", {
      request_id: requestId,
      method: req.method,
      path: req.originalUrl,
      status_code: res.statusCode,
      duration_ms: Date.now() - start,
    });
  });
  next();
});

const chatJsonParser = express.json({ limit: "50mb" });

app.get("/", (req, res) => {
  log("info", "root_status", { request_id: req.requestId });
  res.json({
    status: "ok",
    model: MODEL_ID,
    nvidia_chat_completions_configured: Boolean(NVIDIA_API_KEY),
  });
});

app.get("/health", (req, res) => {
  log("info", "health_checked", {
    request_id: req.requestId,
    model_loaded: model !== null && processor !== null,
  });
  res.json({
    status: "healthy",
    model_loaded: model !== null,
    nvidia_chat_completions_configured: Boolean(NVIDIA_API_KEY),
  });
});

app.post(
  "/v1/chat/completions",
  requireApiKey,
  chatJsonParser,
  async (req, res) => {
    const body = req.body;
    if (!body?.messages || !Array.isArray(body.messages)) {
      return res.status(400).json({
        error: { message: "messages (array) is required." },
      });
    }

    const streamRequested = Boolean(body.stream);
    log("info", "chat_completions_request", {
      request_id: req.requestId,
      stream: streamRequested,
      model: body.model,
      nvidia_configured: Boolean(NVIDIA_API_KEY),
    });

    const proxied = await proxyNvidiaChatCompletions(req, body, streamRequested, res);
    if (proxied || res.headersSent) return;

    log("info", "chat_completions_fallback_local", { request_id: req.requestId });
    return fallbackLocalChatCompletion(req, res, body, streamRequested);
  },
);

app.post("/prompt", requireApiKey, express.json(), async (req, res) => {
  const prompt = req.body.prompt;
  const maxTokens = parseInt(req.body.max_tokens) || 256;
  log("info", "prompt_request_received", {
    request_id: req.requestId,
    prompt_chars: prompt?.length ?? 0,
    max_tokens: maxTokens,
  });

  if (!model || !processor) {
    log("error", "prompt_model_unavailable", { request_id: req.requestId });
    return res.status(503).json({ detail: "Model not loaded yet." });
  }
  if (!prompt) {
    log("error", "prompt_validation_failed", {
      request_id: req.requestId,
      reason: "missing_prompt",
    });
    return res.status(400).json({ detail: "No prompt provided." });
  }

  try {
    await withPromptTrace(req, prompt, maxTokens, async (span) => {
      const generation = span?.startObservation(
        "qwen_text_generation",
        {
          model: MODEL_ID,
          input: prompt,
          modelParameters: { max_new_tokens: maxTokens, do_sample: 0 },
        },
        { asType: "generation" },
      );

      const start = Date.now();
      const response = await queueTextInference(prompt, maxTokens);
      const duration = Date.now() - start;

      generation?.update({ output: response }).end();
      span?.update({
        output: { response_chars: response?.length ?? 0 },
        metadata: { duration_ms: duration },
      });

      log("info", "prompt_completed", {
        request_id: req.requestId,
        duration_ms: duration,
        response_chars: response?.length ?? 0,
      });

      res.json({ response });
    });
  } catch (err) {
    log("error", "prompt_failed", {
      request_id: req.requestId,
      error: err.message,
      stack: err.stack,
    });
    res.status(500).json({ detail: "Inference failed.", error: err.message });
  }
});

Promise.all([setupTracing(), loadModel()]).then(() => {
  app.listen(PORT, "0.0.0.0", async () => {
    log("info", "server_started", {
      host: "0.0.0.0",
      port: PORT,
      model_id: MODEL_ID,
      langfuse_enabled: LANGFUSE_ENABLED,
      nvidia_chat_url: NVIDIA_CHAT_URL,
      nvidia_configured: Boolean(NVIDIA_API_KEY),
    });

    const shutdown = async (signal) => {
      log("info", "shutdown_started", { signal });
      if (telemetrySdk) {
        await telemetrySdk.shutdown();
      }
      process.exit(0);
    };

    process.on("SIGINT", () => shutdown("SIGINT"));
    process.on("SIGTERM", () => shutdown("SIGTERM"));
  });
});