Yuuki-chat/app/api/chat/route.ts

import { NextRequest, NextResponse } from "next/server";

const HF_MODELS: Record<string, string> = {
  "yuuki-v0.1": "YuuKi-OS/Yuuki-v0.1",
  "yuuki-3.7": "YuuKi-OS/Yuuki-3.7",
  "yuuki-best": "YuuKi-OS/Yuuki-best",
};

const YUUKI_API_MODELS: Record<string, string> = {
  "yuuki-v0.1": "yuuki-v0.1",
  "yuuki-3.7": "yuuki-3.7",
  "yuuki-best": "yuuki-best",
};

/**
 * Calls the Yuuki API (yuuki-api.vercel.app) with a yk- token.
 * This is an OpenAI-compatible endpoint.
 */
async function callYuukiApi(
  token: string,
  model: string,
  messages: { role: string; content: string }[]
) {
  const modelId = YUUKI_API_MODELS[model] || "yuuki-best";

  const response = await fetch("https://yuuki-api.vercel.app/api/chat", {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${token}`,
    },
    body: JSON.stringify({
      model: modelId,
      messages,
      max_tokens: 1024,
      temperature: 0.7,
    }),
  });

  if (!response.ok) {
    const errorText = await response.text();
    throw new Error(
      `Yuuki API error (${response.status}): ${errorText.slice(0, 200)}`
    );
  }

  const data = await response.json();
  const content =
    data.choices?.[0]?.message?.content || data.content || "No response";
  return { content, id: data.id || `chatcmpl-${Date.now()}`, model: modelId };
}

/**
 * Calls HuggingFace Inference API via the new router.huggingface.co endpoint.
 * Uses the OpenAI-compatible chat completions format.
 */
async function callHuggingFace(
  token: string,
  model: string,
  messages: { role: string; content: string }[]
) {
  const modelId = HF_MODELS[model] || HF_MODELS["yuuki-best"];
  const url = `https://router.huggingface.co/hf-inference/models/${modelId}/v1/chat/completions`;

  const response = await fetch(url, {
    method: "POST",
    headers: {
      Authorization: `Bearer ${token}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify({
      model: modelId,
      messages,
      max_tokens: 1024,
      temperature: 0.7,
      top_p: 0.9,
    }),
  });

  if (!response.ok) {
    const errorText = await response.text();
    throw new Error(
      `HuggingFace error (${response.status}): ${errorText.slice(0, 200)}`
    );
  }

  const data = await response.json();
  const content =
    data.choices?.[0]?.message?.content?.trim() || "No response generated.";

  return {
    content,
    id: data.id || `chatcmpl-${Date.now()}`,
    model,
  };
}

export async function POST(req: NextRequest) {
  try {
    const body = await req.json();
    const { messages, model, token, tokenSource } = body;

    if (!messages || !Array.isArray(messages)) {
      return NextResponse.json(
        { error: "messages is required" },
        { status: 400 }
      );
    }

    const modelKey = model || "yuuki-best";
    if (!HF_MODELS[modelKey]) {
      return NextResponse.json({ error: "Invalid model" }, { status: 400 });
    }

    let result;

    if (tokenSource === "demo") {
      // Demo mode: use server-side HF_DEMO_TOKEN directly against HuggingFace
      const demoToken = process.env.HF_DEMO_TOKEN;
      if (!demoToken) {
        return NextResponse.json(
          { error: "Demo token not configured on server" },
          { status: 500 }
        );
      }
      result = await callHuggingFace(demoToken, modelKey, messages);
    } else if (tokenSource === "yuuki-api") {
      // Yuuki API: yk- tokens go to yuuki-api.vercel.app
      if (!token) {
        return NextResponse.json(
          { error: "No API token provided" },
          { status: 401 }
        );
      }
      result = await callYuukiApi(token, modelKey, messages);
    } else if (tokenSource === "huggingface") {
      // HuggingFace: hf_ tokens go directly to HF Inference API
      if (!token) {
        return NextResponse.json(
          { error: "No API token provided" },
          { status: 401 }
        );
      }
      result = await callHuggingFace(token, modelKey, messages);
    } else {
      return NextResponse.json(
        { error: "Invalid token source" },
        { status: 400 }
      );
    }

    return NextResponse.json(result);
  } catch (error) {
    const message = error instanceof Error ? error.message : "Unknown error";
    return NextResponse.json({ error: message }, { status: 500 });
  }
}