plates-server/src/ai-coach/ai-coach.controller.ts

import { Body, Controller, Delete, Get, HttpException, HttpStatus, Logger, Param, Post, Query, Res, StreamableFile, UseGuards } from '@nestjs/common';
import { ApiTags, ApiOperation, ApiBody, ApiQuery, ApiParam } from '@nestjs/swagger';
import { Response } from 'express';
import { JwtAuthGuard } from '../common/guards/jwt-auth.guard';
import { CurrentUser } from '../common/decorators/current-user.decorator';
import { AccessTokenPayload } from '../users/services/apple-auth.service';
import { AiCoachService } from './ai-coach.service';
import { AiChatRequestDto, AiChatResponseDto, AiResponseDataDto } from './dto/ai-chat.dto';
import { PostureAssessmentRequestDto, PostureAssessmentResponseDto } from './dto/posture-assessment.dto';
import { UsersService } from '../users/users.service';

@ApiTags('ai-coach')
@Controller('ai-coach')
@UseGuards(JwtAuthGuard)
export class AiCoachController {
  private readonly logger = new Logger(AiCoachController.name);
  constructor(
    private readonly aiCoachService: AiCoachService,
    private readonly usersService: UsersService,
  ) { }

  @Post('chat')
  @ApiOperation({ summary: '流式大模型对话（普拉提教练）' })
  @ApiBody({ type: AiChatRequestDto })
  async chat(
    @Body() body: AiChatRequestDto,
    @CurrentUser() user: AccessTokenPayload,
    @Res({ passthrough: false }) res: Response,
  ): Promise<StreamableFile | AiChatResponseDto | void> {
    const userId = user.sub;
    const stream = body.stream !== false; // 默认流式
    const userContent = body.messages?.[body.messages.length - 1]?.content || '';

    // 判断用户是否有聊天次数
    const usageCount = await this.usersService.getUserUsageCount(userId);
    if (usageCount <= 0) {
      this.logger.error(`chat: ${userId} has no usage count`);
      throw new HttpException('用户没有聊天次数', HttpStatus.FORBIDDEN);
    }

    // 创建或沿用会话ID，并保存用户消息
    const { conversationId } = await this.aiCoachService.createOrAppendMessages({
      userId,
      conversationId: body.conversationId,
      userContent,
    });

    // 体重和饮食指令处理现在已经集成到 streamChat 方法中
    // 通过 # 字符开头的指令系统进行统一处理

    const result = await this.aiCoachService.streamChat({
      userId,
      conversationId,
      userContent,
      imageUrls: body.imageUrls,
      selectedChoiceId: body.selectedChoiceId,
      confirmationData: body.confirmationData,
    });

    await this.usersService.deductUserUsageCount(userId);

    // 检查是否返回结构化数据（如确认选项）
    // 结构化数据必须使用非流式模式返回
    if (typeof result === 'object' && 'type' in result) {
      res.setHeader('Content-Type', 'application/json; charset=utf-8');
      res.send({
        conversationId,
        data: result.data
      });
      return;
    }

    // 普通流式/非流式响应
    const readable = result as any;

    if (!stream) {
      // 非流式：聚合后一次性返回文本
      let text = '';
      for await (const chunk of readable) {
        text += chunk.toString();
      }
      res.setHeader('Content-Type', 'application/json; charset=utf-8');
      res.send({ conversationId, text });
      return;
    }

    // 流式：SSE/文本流
    res.setHeader('Content-Type', 'text/plain; charset=utf-8');
    res.setHeader('Cache-Control', 'no-cache');
    res.setHeader('Transfer-Encoding', 'chunked');

    readable.on('data', (chunk) => {
      res.write(chunk);
    });
    readable.on('end', () => {
      res.end();
    });
    readable.on('error', () => {
      res.end();
    });

    return;
  }

  @Get('conversations')
  @ApiOperation({ summary: '获取会话列表' })
  async listConversations(
    @CurrentUser() user: AccessTokenPayload,
    @Query() query: { page?: number, pageSize?: number },
  ) {
    return this.aiCoachService.listConversations(user.sub, { page: query.page, pageSize: query.pageSize });
  }

  @Get('conversations/:conversationId')
  @ApiOperation({ summary: '获取会话详情（包含消息）' })
  @ApiParam({ name: 'conversationId' })
  async getConversationDetail(
    @CurrentUser() user: AccessTokenPayload,
    @Param('conversationId') conversationId: string,
  ) {
    const data = await this.aiCoachService.getConversationDetail(user.sub, conversationId);
    return data || { code: 404, message: '会话不存在' };
  }

  @Delete('conversations/:conversationId')
  @ApiOperation({ summary: '删除会话（级联删除消息）' })
  @ApiParam({ name: 'conversationId' })
  async deleteConversation(
    @CurrentUser() user: AccessTokenPayload,
    @Param('conversationId') conversationId: string,
  ) {
    const ok = await this.aiCoachService.deleteConversation(user.sub, conversationId);
    return { success: ok };
  }

  @Post('posture-assessment')
  @ApiOperation({ summary: 'AI体态评估' })
  @ApiBody({ type: PostureAssessmentRequestDto })
  async postureAssessment(
    @Body() body: PostureAssessmentRequestDto,
    @CurrentUser() user: AccessTokenPayload,
  ): Promise<PostureAssessmentResponseDto> {
    const res = await this.aiCoachService.assessPosture({
      userId: user.sub,
      frontImageUrl: body.frontImageUrl,
      sideImageUrl: body.sideImageUrl,
      backImageUrl: body.backImageUrl,
      heightCm: body.heightCm,
      weightKg: body.weightKg,
    });
    return res as any;
  }
}