investor-flow/app/server/src/trpc/router.ts

import { initTRPC, TRPCError } from '@trpc/server';
import { z } from 'zod';
import { randomUUID } from 'node:crypto';
import type { Context } from './context.ts';
import { hashPassword, verifyPassword, createSession, clearCookie, oauthStateCookie, verifyOAuthState, OAUTH_STATE_COOKIE } from './context.ts';
import { generateBase32Secret, totp as computeTotp, verifyTotp, otpauthUrl } from '../auth/totp.ts';
import { generateBackupCodes, hashBackupCode } from '../auth/backup-codes.ts';
import { buildAuthorizeUrl, generateState, exchangeCode, type OAuthProvider } from '../auth/oauth.ts';
import type { Quote, PriceCandle, SymbolMeta } from '../cache/CacheRepository.ts';

const t = initTRPC.context<Context>().create();
const router = t.router;
const publicProcedure = t.procedure;

const protectedProcedure = publicProcedure.use(({ ctx, next }) => {
  if (!ctx.userId) throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Sign in required.' });
  return next({ ctx });
});

function oauthCreds(provider: OAuthProvider): { clientId?: string; clientSecret?: string } {
  if (provider === 'github') return { clientId: process.env.GITHUB_CLIENT_ID, clientSecret: process.env.GITHUB_CLIENT_SECRET };
  return { clientId: process.env.GOOGLE_CLIENT_ID, clientSecret: process.env.GOOGLE_CLIENT_SECRET };
}

const authRouter = router({
  signup: publicProcedure
    .input(z.object({ email: z.string().email(), password: z.string().min(8) }))
    .mutation(async ({ ctx, input }) => {
      const email = input.email.toLowerCase();
      const existing = ctx.db.prepare('SELECT id FROM users WHERE email=?').get(email);
      if (existing) throw new TRPCError({ code: 'CONFLICT', message: 'That email is already registered.' });
      const userId = randomUUID();
      ctx.db.prepare('INSERT INTO users (id,email,pw_hash,created_at) VALUES (?,?,?,?)').run(userId, email, hashPassword(input.password), new Date().toISOString());
      const { cookie } = createSession(ctx.db, userId);
      ctx.resHeaders.append('Set-Cookie', cookie);
      return { userId };
    }),
  login: publicProcedure
    .input(z.object({ email: z.string().email(), password: z.string(), totp: z.string().optional() }))
    .mutation(async ({ ctx, input }) => {
      const email = input.email.toLowerCase();
      const row = ctx.db.prepare('SELECT id, pw_hash, is_2fa_enabled, totp_secret FROM users WHERE email=?').get(email) as { id: string; pw_hash: string; is_2fa_enabled: number; totp_secret: string | null } | undefined;
      if (!row || !verifyPassword(input.password, row.pw_hash)) throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Invalid email or password.' });
      if (row.is_2fa_enabled === 1) {
        if (!input.totp || !row.totp_secret || !verifyTotp(input.totp, row.totp_secret)) {
          throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Two-factor code required or invalid.' });
        }
      }
      const { cookie } = createSession(ctx.db, row.id);
      ctx.resHeaders.append('Set-Cookie', cookie);
      return { userId: row.id };
    }),
  logout: publicProcedure.mutation(({ ctx }) => {
    ctx.resHeaders.append('Set-Cookie', clearCookie());
    return { ok: true };
  }),
  me: publicProcedure.query(({ ctx }) => {
    if (!ctx.userId) return null;
    const u = ctx.db.prepare('SELECT id,email,complexity,risk_tolerance,convexity_posture FROM users WHERE id=?').get(ctx.userId) as { id: string; email: string; complexity: string; risk_tolerance: string; convexity_posture: string } | undefined;
    return u ? { userId: u.id, email: u.email, complexity: u.complexity, riskTolerance: u.risk_tolerance, convexityPosture: u.convexity_posture } : null;
  }),
  enable2fa: protectedProcedure.input(z.object({})).mutation(async ({ ctx }) => {
    const userId = ctx.userId as string;
    const u = ctx.db.prepare('SELECT email FROM users WHERE id=?').get(userId) as { email: string } | undefined;
    const secret = generateBase32Secret();
    const codes = generateBackupCodes(10);
    const hashes = codes.map((c) => hashBackupCode(c));
    ctx.db.prepare('UPDATE users SET totp_secret=?, backup_codes_hashed=? WHERE id=?').run(secret, JSON.stringify(hashes), userId);
    return { totpSecret: secret, qrUrl: otpauthUrl(secret, 'Investor Flow', u?.email ?? 'user'), backupCodes: codes };
  }),
  confirm2fa: protectedProcedure.input(z.object({ totp: z.string() })).mutation(async ({ ctx, input }) => {
    const userId = ctx.userId as string;
    const row = ctx.db.prepare('SELECT totp_secret FROM users WHERE id=?').get(userId) as { totp_secret: string | null } | undefined;
    if (!row?.totp_secret) throw new TRPCError({ code: 'BAD_REQUEST', message: 'Enable two-factor first.' });
    if (!verifyTotp(input.totp, row.totp_secret)) throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Invalid two-factor code.' });
    ctx.db.prepare('UPDATE users SET is_2fa_enabled=1 WHERE id=?').run(userId);
    return { ok: true };
  }),
  // Slice 2b — OAuth start: return the provider authorize URL + set a CSRF state cookie.
  oauthStart: publicProcedure
    .input(z.object({ provider: z.enum(['github', 'google']), redirectUri: z.string().url() }))
    .mutation(({ ctx, input }) => {
      const creds = oauthCreds(input.provider);
      if (!creds.clientId || !creds.clientSecret) throw new TRPCError({ code: 'BAD_REQUEST', message: 'OAuth provider not configured.' });
      const state = generateState();
      ctx.resHeaders.append('Set-Cookie', oauthStateCookie(input.provider, state, input.redirectUri));
      return { redirectUrl: buildAuthorizeUrl(input.provider, { clientId: creds.clientId, redirectUri: input.redirectUri, state }), state };
    }),
  // Slice 2b — OAuth callback: verify CSRF state, exchange code, find/link/create user, start a session.
  oauthCallback: publicProcedure
    .input(z.object({ provider: z.enum(['github', 'google']), code: z.string(), state: z.string(), redirectUri: z.string().url() }))
    .mutation(async ({ ctx, input }) => {
      if (!verifyOAuthState(ctx.cookies[OAUTH_STATE_COOKIE], input.provider, input.state, input.redirectUri)) {
        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Invalid OAuth state.' });
      }
      const creds = oauthCreds(input.provider);
      if (!creds.clientId || !creds.clientSecret) throw new TRPCError({ code: 'BAD_REQUEST', message: 'OAuth provider not configured.' });
      const info = await exchangeCode(input.provider, input.code, { clientId: creds.clientId, clientSecret: creds.clientSecret, redirectUri: input.redirectUri });
      const subject = info.providerSubject;
      const email = info.email.toLowerCase();
      if (!subject || !email) throw new TRPCError({ code: 'BAD_REQUEST', message: 'Provider did not return a usable identity.' });
      // existing link?
      let row = ctx.db.prepare('SELECT id FROM users WHERE oauth_subject=? AND oauth_provider=?').get(subject, input.provider) as { id: string } | undefined;
      if (!row) {
        // link existing account by email, else create an OAuth-only account (sentinel pw_hash)
        const byEmail = ctx.db.prepare('SELECT id FROM users WHERE email=?').get(email) as { id: string } | undefined;
        const userId = byEmail?.id ?? randomUUID();
        if (byEmail) {
          ctx.db.prepare('UPDATE users SET oauth_subject=?, oauth_provider=? WHERE id=?').run(subject, input.provider, userId);
        } else {
          ctx.db.prepare('INSERT INTO users (id,email,pw_hash,oauth_subject,oauth_provider,created_at) VALUES (?,?,?,?,?,?)').run(userId, email, 'oauth', subject, input.provider, new Date().toISOString());
        }
        row = { id: userId };
      }
      const { cookie } = createSession(ctx.db, row.id);
      ctx.resHeaders.append('Set-Cookie', cookie);
      return { userId: row.id };
    }),
});

const marketRouter = router({
  snapshot: publicProcedure
    .input(z.object({ symbol: z.string().min(1) }))
    .query(async ({ ctx, input }) => {
      const symbol = input.symbol.toUpperCase();
      const k = { quote: `yfinance:quote:${symbol}`, candles: `yfinance:candles:${symbol}:1d`, sector: `yfinance:symbol:${symbol}` };
      const entries = await ctx.cache.getMany<unknown>([k.quote, k.candles, k.sector]);
      const byKey = new Map(entries.map((e) => [e.key, e]));
      const val = <T>(key: string): T | null => (byKey.get(key)?.value ?? null) as T | null;
      const stale = (key: string): boolean => byKey.get(key)?.isStale ?? true;
      return {
        symbol,
        quote: val<Quote>(k.quote),
        candles: val<PriceCandle[]>(k.candles),
        sector: val<SymbolMeta>(k.sector),
        stale: { quote: stale(k.quote), candles: stale(k.candles), sector: stale(k.sector) },
      };
    }),
});

export const appRouter = router({ auth: authRouter, market: marketRouter });
export type AppRouter = typeof appRouter;