import { NextRequest } from "next/server";
import { db } from "@/lib/db";
import { resolveApiKeyUser } from "@/lib/api-auth";

/**
 * POST /api/relay/v1/chat/completions
 *
 * The GPU-as-API relay. This is the core innovation:
 * 1. Validate API key → resolve to user
 * 2. Find user's online device(s)
 * 3. Forward request to device's tunnel URL
 * 4. Stream response back to the caller
 *
 * The caller uses this as a drop-in OpenAI replacement:
 *   OPENAI_BASE_URL=https://llmhosts.com/api/relay/v1
 */
export async function POST(req: NextRequest) {
  // 1. Authenticate
  const user = await resolveApiKeyUser(req);
  if (!user) {
    return new Response(
      JSON.stringify({ error: { message: "Invalid API key", type: "authentication_error" } }),
      { status: 401, headers: { "Content-Type": "application/json" } },
    );
  }

  // 2. Find user's online devices
  const devices = await db.device.findMany({
    where: {
      userId: user.id,
      status: "online",
    },
    orderBy: { lastSeen: "desc" },
  });

  if (devices.length === 0) {
    return new Response(
      JSON.stringify({
        error: {
          message: "No online devices. Run `llmhosts up` on a machine with a GPU to connect it.",
          type: "device_error",
        },
      }),
      { status: 503, headers: { "Content-Type": "application/json" } },
    );
  }

  // 3. Parse request body to check model
  let body: Record<string, unknown>;
  try {
    body = await req.json();
  } catch {
    return new Response(
      JSON.stringify({ error: { message: "Invalid JSON body", type: "invalid_request_error" } }),
      { status: 400, headers: { "Content-Type": "application/json" } },
    );
  }

  const requestedModel = (body.model as string) || "";

  // 4. Select best device (model match + lowest latency)
  let targetDevice = devices[0]; // Default: most recently seen

  if (requestedModel) {
    for (const device of devices) {
      try {
        const models: string[] = JSON.parse(device.models);
        if (models.some((m) => m === requestedModel || m.startsWith(requestedModel.split(":")[0]))) {
          targetDevice = device;
          break;
        }
      } catch {
        continue;
      }
    }
  }

  // 5. Check device freshness (offline if last seen > 90s ago)
  const staleness = Date.now() - targetDevice.lastSeen.getTime();
  if (staleness > 90_000) {
    // Mark stale and reject
    await db.device.update({
      where: { id: targetDevice.id },
      data: { status: "offline" },
    });
    return new Response(
      JSON.stringify({
        error: {
          message: "Device appears offline (no heartbeat in 90s). Check your `llmhosts up` session.",
          type: "device_error",
        },
      }),
      { status: 503, headers: { "Content-Type": "application/json" } },
    );
  }

  // 6. Forward to device via tunnel
  const tunnelUrl = targetDevice.tunnelUrl.replace(/\/$/, "");
  const targetUrl = `${tunnelUrl}/v1/chat/completions`;

  const isStreaming = body.stream === true;

  try {
    const upstreamResponse = await fetch(targetUrl, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "X-LLMHosts-Relay": "true",
        "X-LLMHosts-User": user.id,
      },
      body: JSON.stringify(body),
      signal: AbortSignal.timeout(120_000), // 2 minute timeout
    });

    if (!upstreamResponse.ok) {
      const errorBody = await upstreamResponse.text();
      return new Response(errorBody, {
        status: upstreamResponse.status,
        headers: {
          "Content-Type": upstreamResponse.headers.get("Content-Type") || "application/json",
          "X-LLMHosts-Device": targetDevice.name,
        },
      });
    }

    // 7. Stream response back (SSE or JSON)
    if (isStreaming && upstreamResponse.body) {
      return new Response(upstreamResponse.body, {
        status: 200,
        headers: {
          "Content-Type": "text/event-stream",
          "Cache-Control": "no-cache",
          Connection: "keep-alive",
          "X-LLMHosts-Device": targetDevice.name,
          "X-LLMHosts-Relay": "true",
        },
      });
    }

    // Non-streaming: return JSON response
    const responseBody = await upstreamResponse.text();
    return new Response(responseBody, {
      status: 200,
      headers: {
        "Content-Type": "application/json",
        "X-LLMHosts-Device": targetDevice.name,
        "X-LLMHosts-Relay": "true",
      },
    });
  } catch (err) {
    const message = err instanceof Error ? err.message : "Unknown error";
    return new Response(
      JSON.stringify({
        error: {
          message: `Failed to reach device "${targetDevice.name}": ${message}`,
          type: "relay_error",
        },
      }),
      { status: 502, headers: { "Content-Type": "application/json" } },
    );
  }
}
