import { Router, Request } from 'express';
import jwt from 'jsonwebtoken';
import crypto from 'crypto';
import { getTunnel, TunnelConfig } from '../db';
import { registry } from '../relay/TunnelRegistry';
import { AuthPayload, extractBearer } from '../middleware/requireAuth';
import { addPending, removePending, RelayResponse } from '../relay/pending';

export interface RelayRequest
{
  reqId: string;
  method: string;
  path: string;
  headers: Record<string, string>;
  body: string; // base64-encoded
}

const router = Router();

function softAuth( req: Request ): void
{
  const token = req.cookies?.accessToken ?? extractBearer( req );
  if ( !token ) return;
  try
  {
    req.auth = jwt.verify( token, process.env.JWT_SECRET ?? '' ) as AuthPayload;
  }
  catch { /* leave req.auth undefined */ }
}

function canAccess( tunnel: TunnelConfig, req: Request ): boolean
{
  if ( tunnel.access === 'public' ) return true;
  if ( !req.auth ) return false;
  return tunnel.ownerId === req.auth.userId || tunnel.allowedUserIds.includes( req.auth.userId );
}

// Matches /:tunnelId and /:tunnelId/any/path
router.all( [ '/:tunnelId', '/:tunnelId/*' ], async ( req, res ) =>
{
  softAuth( req );

  const tunnel = getTunnel( req.params.tunnelId );
  if ( !tunnel ) { res.status( 404 ).json( { error: 'Tunnel not found' } ); return; }
  if ( !canAccess( tunnel, req ) ) { res.status( 401 ).json( { error: 'Not authenticated' } ); return; }

  const ws = registry.get( tunnel.id );
  if ( !ws ) { res.status( 503 ).json( { error: 'Tunnel agent not connected' } ); return; }

  // Strip the /:tunnelId prefix to get the path to forward
  const prefix = '/' + req.params.tunnelId;
  const forwardPath = req.url.slice( prefix.length ) || '/';

  // Buffer the raw request body (no JSON middleware on /t routes)
  const chunks: Buffer[] = [];
  await new Promise<void>( ( resolve, reject ) =>
  {
    req.on( 'data', chunk => chunks.push( Buffer.from( chunk ) ) );
    req.on( 'end', resolve );
    req.on( 'error', reject );
  } );
  const bodyBuf = Buffer.concat( chunks );

  // Forward all headers except host and content-length
  const headers: Record<string, string> = {};
  for ( const [ key, val ] of Object.entries( req.headers ) )
  {
    if ( key === 'host' || key === 'content-length' ) continue;
    if ( typeof val === 'string' ) headers[ key ] = val;
    else if ( Array.isArray( val ) ) headers[ key ] = val.join( ', ' );
  }

  const reqId = crypto.randomUUID();

  const relayReq: RelayRequest = {
    reqId,
    method: req.method,
    path: forwardPath,
    headers,
    body: bodyBuf.toString( 'base64' ),
  };

  const TIMEOUT_MS = 30_000;
  const response = await new Promise<RelayResponse | null>( ( resolve ) =>
  {
    const timer = setTimeout( () =>
    {
      removePending( reqId );
      resolve( null );
    }, TIMEOUT_MS );

    addPending( reqId, ( resp ) =>
    {
      clearTimeout( timer );
      resolve( resp );
    } );

    ws.send( JSON.stringify( relayReq ) );
  } );

  if ( !response ) { res.status( 504 ).json( { error: 'Agent timed out' } ); return; }

  res.status( response.status );
  for ( const [ key, val ] of Object.entries( response.headers ) )
  {
    if ( key.toLowerCase() === 'transfer-encoding' ) continue;
    res.setHeader( key, val );
  }
  res.end( Buffer.from( response.body, 'base64' ) );
} );

export default router;