ai-learning-platform/backend/src/modules/sandbox/sandbox.service.ts

import { Injectable, HttpException, HttpStatus } from '@nestjs/common';
import { PrismaService } from '../../prisma/prisma.service';
import { AIGatewayService, ChatOptions } from '../ai/ai-gateway.service';
import { randomUUID, createHmac } from 'crypto';

export interface SandboxModelInfo {
  id: string;
  label: string;
  provider: string;
  desc: string;
}

interface StreamResult {
  type: 'text' | 'done' | 'error';
  content?: string;
  sessionId?: number;
  conversationId?: string;
  message?: string;
}

@Injectable()
export class SandboxService {
  constructor(
    private prisma: PrismaService,
    private aiGateway: AIGatewayService,
  ) {}

  async chat(userId: number, conversationId: string | undefined, model: string, messages: { role: string; content: string }[], options?: ChatOptions, images?: string[]) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user || user.status !== 'ACTIVE') {
      throw new HttpException('用户不可用', HttpStatus.FORBIDDEN);
    }

    const convId = conversationId || randomUUID();

    // 配额检查：先用每日免费配额，再用购买的额外配额
    const today = new Date();
    today.setHours(0, 0, 0, 0);
    const existing = await this.prisma.sandboxSession.findUnique({
      where: { userId_conversationId: { userId, conversationId: convId } },
    });
    if (!existing) {
      const todayCount = await this.prisma.sandboxSession.count({
        where: { userId, createdAt: { gte: today } },
      });
      const dailyLimit = user.sandboxDaily || 10;
      if (todayCount >= dailyLimit) {
        if ((user.sandboxExtra || 0) > 0) {
          await this.prisma.user.update({
            where: { id: userId },
            data: { sandboxExtra: { decrement: 1 } },
          });
        } else {
          throw new HttpException('今日沙箱使用次数已用完', HttpStatus.TOO_MANY_REQUESTS);
        }
      }
    }

    // Convert images to structured message content
    let aiMessages = messages as any[];
    if (images && images.length > 0) {
      aiMessages = messages.map(m => {
        if (m.role === 'user' && m === messages[messages.length - 1]) {
          const parts: any[] = [{ type: 'text', text: m.content }];
          for (const img of images) {
            parts.push({ type: 'image_url', image_url: { url: img } });
          }
          return { role: m.role, content: parts };
        }
        return m;
      });
    }

    let reply = await this.aiGateway.chat(model, aiMessages, options);
    if (typeof reply !== 'string') {
      reply = '抱歉，AI 返回了无效的回复，请重试。';
    }

    const allMessages = messages.concat({ role: 'assistant', content: reply });
    const firstUserMsg = messages.find(m => m.role === 'user');
    const title = firstUserMsg ? firstUserMsg.content.slice(0, 80) : 'AI 对话';

    const session = await this.prisma.sandboxSession.upsert({
      where: { userId_conversationId: { userId, conversationId: convId } },
      create: {
        userId,
        conversationId: convId,
        model,
        title,
        messages: JSON.stringify(allMessages),
        tokens: Math.ceil(reply.length / 2),
      },
      update: {
        model,
        title,
        messages: JSON.stringify(allMessages),
        tokens: Math.ceil(reply.length / 2),
      },
    });

    return { reply, conversationId: convId, sessionId: session.id };
  }

  async *chatStream(userId: number, conversationId: string | undefined, model: string, messages: { role: string; content: string }[], options?: ChatOptions, images?: string[]): AsyncGenerator<string> {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user || user.status !== 'ACTIVE') {
      yield JSON.stringify({ type: 'error', message: '用户不可用' } as StreamResult);
      return;
    }

    const convId = conversationId || randomUUID();

    const today = new Date();
    today.setHours(0, 0, 0, 0);
    const existing = await this.prisma.sandboxSession.findUnique({
      where: { userId_conversationId: { userId, conversationId: convId } },
    });
    if (!existing) {
      const todayCount = await this.prisma.sandboxSession.count({
        where: { userId, createdAt: { gte: today } },
      });
      const dailyLimit = user.sandboxDaily || 10;
      if (todayCount >= dailyLimit) {
        if ((user.sandboxExtra || 0) > 0) {
          await this.prisma.user.update({
            where: { id: userId },
            data: { sandboxExtra: { decrement: 1 } },
          });
        } else {
          yield JSON.stringify({ type: 'error', message: '今日沙箱使用次数已用完' } as StreamResult);
          return;
        }
      }
    }

    let aiMessages = messages as any[];
    if (images && images.length > 0) {
      aiMessages = messages.map(m => {
        if (m.role === 'user' && m === messages[messages.length - 1]) {
          const parts: any[] = [{ type: 'text', text: m.content }];
          for (const img of images) {
            parts.push({ type: 'image_url', image_url: { url: img } });
          }
          return { role: m.role, content: parts };
        }
        return m;
      });
    }

    let fullReply = '';
    try {
      for await (const chunk of this.aiGateway.chatStream(model, aiMessages, options)) {
        fullReply += chunk;
        yield JSON.stringify({ type: 'text', content: chunk } as StreamResult);
      }
    } catch (err: any) {
      yield JSON.stringify({ type: 'error', message: err.message } as StreamResult);
      return;
    }

    if (typeof fullReply !== 'string' || fullReply.length === 0) {
      fullReply = '抱歉，AI 返回了无效的回复，请重试。';
    }

    const allMessages = messages.concat({ role: 'assistant', content: fullReply });
    const firstUserMsg = messages.find(m => m.role === 'user');
    const title = firstUserMsg ? firstUserMsg.content.slice(0, 80) : 'AI 对话';

    const session = await this.prisma.sandboxSession.upsert({
      where: { userId_conversationId: { userId, conversationId: convId } },
      create: {
        userId,
        conversationId: convId,
        model,
        title,
        messages: JSON.stringify(allMessages),
        tokens: Math.ceil(fullReply.length / 2),
      },
      update: {
        model,
        title,
        messages: JSON.stringify(allMessages),
        tokens: Math.ceil(fullReply.length / 2),
      },
    });

    yield JSON.stringify({ type: 'done', sessionId: session.id, conversationId: convId } as StreamResult);
  }

  async getSessions(userId: number, params: { page?: number; pageSize?: number; search?: string }) {
    const page = Number(params.page ?? 1);
    const pageSize = Number(params.pageSize ?? 50);
    const where: any = { userId };
    if (params.search) {
      where.title = { contains: params.search };
    }
    const [items, total] = await Promise.all([
      this.prisma.sandboxSession.findMany({
        where,
        skip: (page - 1) * pageSize,
        take: pageSize,
        orderBy: { createdAt: 'desc' },
        select: { id: true, conversationId: true, model: true, title: true, feedback: true, createdAt: true, tokens: true },
      }),
      this.prisma.sandboxSession.count({ where }),
    ]);
    return { items, total, page, pageSize };
  }

  async getSession(userId: number, id: number) {
    const session = await this.prisma.sandboxSession.findFirst({
      where: { id, userId },
    });
    if (!session) {
      throw new HttpException('会话不存在', HttpStatus.NOT_FOUND);
    }
    return {
      id: session.id,
      conversationId: session.conversationId,
      model: session.model,
      title: session.title,
      feedback: session.feedback,
      createdAt: session.createdAt,
      messages: JSON.parse(session.messages),
    };
  }

  async setFeedback(userId: number, id: number, feedback: string | null) {
    const session = await this.prisma.sandboxSession.findFirst({
      where: { id, userId },
    });
    if (!session) {
      throw new HttpException('会话不存在', HttpStatus.NOT_FOUND);
    }
    await this.prisma.sandboxSession.update({
      where: { id },
      data: { feedback: feedback || null },
    });
    return { success: true };
  }

  async deleteSession(userId: number, id: number) {
    const session = await this.prisma.sandboxSession.findFirst({
      where: { id, userId },
    });
    if (!session) {
      throw new HttpException('会话不存在', HttpStatus.NOT_FOUND);
    }
    await this.prisma.sandboxSession.delete({ where: { id } });
    return { success: true };
  }

  async renameSession(userId: number, id: number, title: string) {
    const session = await this.prisma.sandboxSession.findFirst({
      where: { id, userId },
    });
    if (!session) throw new HttpException('会话不存在', HttpStatus.NOT_FOUND);
    await this.prisma.sandboxSession.update({ where: { id }, data: { title } });
    return { success: true };
  }

  async getHistory(userId: number, params: { page?: number; pageSize?: number }) {
    const page = Number(params.page ?? 1);
    const pageSize = Number(params.pageSize ?? 20);
    const [items, total] = await Promise.all([
      this.prisma.sandboxSession.findMany({
        where: { userId },
        skip: (page - 1) * pageSize,
        take: pageSize,
        orderBy: { createdAt: 'desc' },
        select: { id: true, conversationId: true, model: true, title: true, feedback: true, createdAt: true, tokens: true },
      }),
      this.prisma.sandboxSession.count({ where: { userId } }),
    ]);
    return { items, total, page, pageSize };
  }

  async getQuota(userId: number) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    const today = new Date();
    today.setHours(0, 0, 0, 0);
    const used = await this.prisma.sandboxSession.count({
      where: { userId, createdAt: { gte: today } },
    });
    const dailyLimit = user?.sandboxDaily || 10;
    const extra = user?.sandboxExtra || 0;
    const dailyRemaining = Math.max(0, dailyLimit - used);
    return {
      dailyLimit,
      used,
      dailyRemaining,
      extra,
      canUse: dailyRemaining > 0 || extra > 0,
      totalRemaining: dailyRemaining + extra,
    };
  }

  /** 返回沙箱中可用的模型列表 */
  getAvailableModels(): SandboxModelInfo[] {
    return this.aiGateway.getSandboxModels().map(m => ({
      id: m.id,
      label: m.label,
      provider: m.provider,
      desc: m.desc,
    }));
  }

  async generateShareToken(userId: number, sessionId: number) {
    const session = await this.prisma.sandboxSession.findFirst({
      where: { id: sessionId, userId },
    });
    if (!session) throw new HttpException('会话不存在', HttpStatus.NOT_FOUND);

    const secret = process.env.JWT_SECRET || 'yuzhiran-share-secret';
    const ts = Date.now().toString(36);
    const hmac = createHmac('sha256', secret).update(`${sessionId}:${ts}`).digest('hex').slice(0, 12);
    const token = Buffer.from(`${sessionId}|${ts}|${hmac}`).toString('base64url');
    return { shareUrl: `${process.env.FRONTEND_URL || 'http://localhost:3000'}/sandbox/share?token=${token}` };
  }

  async getSessionByShareToken(token: string) {
    try {
      const decoded = Buffer.from(token, 'base64url').toString();
      const parts = decoded.split('|');
      if (parts.length !== 3) throw new Error('invalid token');
      const [sessionId, ts, hmac] = parts;
      const secret = process.env.JWT_SECRET || 'yuzhiran-share-secret';
      const expected = createHmac('sha256', secret).update(`${sessionId}:${ts}`).digest('hex').slice(0, 12);
      if (hmac !== expected) throw new Error('invalid signature');

      const session = await this.prisma.sandboxSession.findUnique({
        where: { id: parseInt(sessionId) },
      });
      if (!session) throw new Error('session not found');

      return {
        id: session.id,
        model: session.model,
        title: session.title,
        createdAt: session.createdAt,
        messages: JSON.parse(session.messages).filter((m: any) => m.role !== 'system'),
      };
    } catch {
      throw new HttpException('分享链接无效或已过期', HttpStatus.BAD_REQUEST);
    }
  }
}