import type { Handle } from '@sveltejs/kit';
import { jwtVerify, createRemoteJWKSet } from 'jose';

const JWT_COOKIE_NAME = 'sb_jwt';
const JWKS_CACHE = new Map<string, ReturnType<typeof createRemoteJWKSet>>();

function getJwks(pdsUrl: string) {
  if (!JWKS_CACHE.has(pdsUrl)) {
    JWKS_CACHE.set(pdsUrl, createRemoteJWKSet(new URL(`${pdsUrl}/.well-known/jwks.json`)));
  }
  return JWKS_CACHE.get(pdsUrl)!;
}

async function validateJwt(token: string, pdsUrl: string): Promise<{
  valid: boolean;
  payload?: {
    sub: string;
    player_id?: string;
    developer_id?: string;
    role?: string;
  };
}> {
  try {
    const JWKS = getJwks(pdsUrl);
    const { payload } = await jwtVerify(token, JWKS, {
      issuer: pdsUrl,
    });
    return { valid: true, payload: payload as { sub: string; player_id?: string; developer_id?: string; role?: string } };
  } catch {
    return { valid: false };
  }
}

function parseCookies(cookieHeader: string | null): Record<string, string> {
  if (!cookieHeader) return {};
  return Object.fromEntries(
    cookieHeader.split(';').map(c => {
      const [key, ...val] = c.trim().split('=');
      return [key, val.join('=')];
    })
  );
}

export const handle: Handle = async ({ event, resolve }) => {
  if (event.url.pathname.startsWith('/api/')) {
    const cookieHeader = event.request.headers.get('cookie');
    const cookies = parseCookies(cookieHeader);
    const token = cookies[JWT_COOKIE_NAME];

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
    };

    if (token) {
      const [headerPart] = token.split('.');
      let pdsUrl = 'https://bsky.social';

      try {
        const headerJson = JSON.parse(atob(headerPart));
        if (headerJson.iss) {
          pdsUrl = headerJson.iss.replace(/\/$/, '');
        }
      } catch {
      }

      const validation = await validateJwt(token, pdsUrl);

      if (validation.valid && validation.payload) {
        headers['Authorization'] = `Bearer ${token}`;
        event.locals.gameDeveloper = validation.payload.developer_id ?? null;
        event.locals.playerId = validation.payload.player_id ?? null;
        event.locals.did = validation.payload.sub;
        event.locals.role = validation.payload.role ?? null;
      } else {
        const response = new Response(JSON.stringify({ error: 'Invalid or expired session' }), {
          status: 401,
          headers: { 'Content-Type': 'application/json' },
        });
        response.headers.append(
          'Set-Cookie',
          `${JWT_COOKIE_NAME}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0`
        );
        return response;
      }
    }

    const targetUrl = `${event.url.pathname}${event.url.search}`;
    const body = event.request.body ? await event.request.clone().arrayBuffer() : undefined;
    const backendResponse = await fetch(targetUrl, {
      method: event.request.method,
      headers,
      body,
      redirect: 'follow',
    });

    const responseBody = await backendResponse.text();
    const responseHeaders = new Headers();
    responseHeaders.set('Content-Type', backendResponse.headers.get('Content-Type') || 'application/json');

    if (backendResponse.status === 401 && token) {
      responseHeaders.append(
        'Set-Cookie',
        `${JWT_COOKIE_NAME}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0`
      );
    }

    return new Response(responseBody, {
      status: backendResponse.status,
      headers: responseHeaders,
    });
  }

  const response = await resolve(event);
  return response;
};