import { Logger } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { OnGatewayConnection, OnGatewayDisconnect, WebSocketGateway } from '@nestjs/websockets';
import type { IncomingMessage } from 'http';
import type { WebSocket } from 'ws';
import { JwtPayload, AUTH_COOKIE } from '../auth/strategies/jwt.strategy';
import { PrismaService } from '../prisma/prisma.service';
import { RealtimeService } from './realtime.service';

interface ClientMessage {
  type: 'subscribe' | 'unsubscribe';
  channel: string;
}

@WebSocketGateway({ path: '/api/ws' })
export class RealtimeGateway implements OnGatewayConnection, OnGatewayDisconnect {
  private readonly logger = new Logger(RealtimeGateway.name);
  private readonly users = new WeakMap<WebSocket, string>();

  constructor(
    private readonly jwtService: JwtService,
    private readonly prisma: PrismaService,
    private readonly realtime: RealtimeService,
  ) {}

  handleConnection(client: WebSocket, req: IncomingMessage) {
    const userId = this.authenticate(req);
    if (!userId) {
      client.close(4401, 'unauthorized');
      return;
    }
    this.users.set(client, userId);
    client.send(JSON.stringify({ type: 'connected' }));

    client.on('message', (data) => {
      this.onClientMessage(client, data.toString()).catch((err) =>
        this.logger.error(`WS message error: ${err.message}`),
      );
    });
  }

  handleDisconnect(client: WebSocket) {
    this.realtime.removeSocket(client);
  }

  private authenticate(req: IncomingMessage): string | null {
    // Same-origin httpOnly cookie rides along on the WS upgrade request
    const cookies = req.headers.cookie ?? '';
    const token = cookies
      .split(';')
      .map((c) => c.trim())
      .find((c) => c.startsWith(`${AUTH_COOKIE}=`))
      ?.slice(AUTH_COOKIE.length + 1);
    if (!token) {
      return null;
    }
    try {
      const payload = this.jwtService.verify<JwtPayload>(token);
      return payload.sub;
    } catch {
      return null;
    }
  }

  private async onClientMessage(client: WebSocket, raw: string) {
    let msg: ClientMessage;
    try {
      msg = JSON.parse(raw);
    } catch {
      client.send(JSON.stringify({ type: 'error', reason: 'invalid JSON' }));
      return;
    }

    if (msg.type === 'unsubscribe' && typeof msg.channel === 'string') {
      this.realtime.unsubscribe(msg.channel, client);
      client.send(JSON.stringify({ type: 'unsubscribed', channel: msg.channel }));
      return;
    }

    if (msg.type !== 'subscribe' || typeof msg.channel !== 'string') {
      client.send(JSON.stringify({ type: 'error', reason: 'expected {type: subscribe|unsubscribe, channel}' }));
      return;
    }

    const allowed = await this.canAccessChannel(this.users.get(client), msg.channel);
    if (!allowed) {
      client.send(JSON.stringify({ type: 'error', reason: `cannot subscribe to ${msg.channel}` }));
      return;
    }

    this.realtime.subscribe(msg.channel, client);
    client.send(JSON.stringify({ type: 'subscribed', channel: msg.channel }));
  }

  private async canAccessChannel(userId: string | undefined, channel: string): Promise<boolean> {
    if (!userId) {
      return false;
    }
    if (channel.startsWith('job:')) {
      const jobId = channel.slice(4);
      const job = await this.prisma.job.findUnique({ where: { id: jobId }, select: { orgId: true } });
      return job ? this.isMember(job.orgId, userId) : false;
    }
    const devicesMatch = /^org:(.+):devices$/.exec(channel);
    if (devicesMatch) {
      return this.isMember(devicesMatch[1], userId);
    }
    return false;
  }

  private async isMember(orgId: string, userId: string): Promise<boolean> {
    const membership = await this.prisma.orgMembership.findUnique({
      where: { orgId_userId: { orgId, userId } },
    });
    return membership !== null;
  }
}
