import type { inferAsyncReturnType } from "@trpc/server";
import type { FetchCreateContextFnOptions } from "@trpc/server/adapters/fetch";
import { db } from "~/server/db";
import { verifyJWT } from "~/server/auth/jwt";
import { validateSession } from "~/server/auth/session";
import { users } from "~/server/db/schema/auth";
import { eq } from "drizzle-orm";

export type CreateTRPCContextOptions = {
  req: Request;
  resHeaders?: Headers;
};

function parseCookies(req: Request): Record<string, string> {
  const cookieHeader = req.headers.get("cookie") ?? "";
  const cookies: Record<string, string> = {};
  for (const cookie of cookieHeader.split(";")) {
    const trimmed = cookie.trim();
    if (!trimmed) continue;
    const idx = trimmed.indexOf("=");
    if (idx === -1) {
      cookies[trimmed] = "";
    } else {
      cookies[trimmed.slice(0, idx).trim()] = trimmed.slice(idx + 1).trim();
    }
  }
  return cookies;
}

export async function createTRPCContext(
  opts: CreateTRPCContextOptions,
): Promise<{
  db: typeof db;
  user: typeof users.$inferSelect | null;
  apiKey: string | null;
}> {
  const { req } = opts;
  let userId: string | null = null;
  let apiKey: string | null = null;

  const cookies = parseCookies(req);
  const sessionToken = cookies["session_token"];

  if (sessionToken) {
    const result = await validateSession(sessionToken);
    if (result) {
      userId = result.user.id;
    }
  }

  if (!userId) {
    const authHeader = req.headers.get("authorization");
    if (authHeader?.startsWith("Bearer ")) {
      const token = authHeader.slice(7);
      try {
        const payload = await verifyJWT<{ sub?: string }>(token);
        userId = payload.sub ?? null;
      } catch {
        // Invalid token
      }
    }
  }

  if (!userId) {
    apiKey = req.headers.get("x-api-key") ?? null;
  }

  let user: typeof users.$inferSelect | null = null;
  if (userId) {
    const [found] = await db
      .select()
      .from(users)
      .where(eq(users.id, userId))
      .limit(1);
    user = found ?? null;
  }

  return { db, user, apiKey };
}

export type TRPCContext = inferAsyncReturnType<typeof createTRPCContext>;
