server/routers/chat.ts

import { z } from "zod";
import { protectedProcedure, router } from "../_core/trpc";
import { invokeLLM } from "../_core/llm";
import { updateProjectStatus, addIteration } from "../db";

export const chatRouter = router({
  sendMessage: protectedProcedure
    .input(
      z.object({
        projectId: z.number().optional(),
        message: z.string().min(1),
        mode: z.enum(["qwen", "deepseek", "loop", "auto"]).default("auto"),
        contentType: z.enum(["code", "exploit", "payload", "information", "strategy"]).default("code"),
      })
    )
    .mutation(async ({ ctx, input }) => {
      try {
        // Call LLM based on mode
        let response;

        if (input.mode === "qwen" || input.mode === "auto") {
          response = await invokeLLM({
            model: "huihui-ai/Qwen2.5-72B-Instruct-abliterated",
            messages: [
              {
                role: "system",
                content: `You are an expert code generator specializing in ${input.contentType} generation. 
Generate high-quality, production-ready ${input.contentType}. 
Focus on: security, efficiency, and reliability.`,
              },
              {
                role: "user",
                content: input.message,
              },
            ],
          });
        } else if (input.mode === "deepseek") {
          response = await invokeLLM({
            model: "huihui-ai/DeepSeek-R1-Distill-Qwen-32B-abliterated",
            messages: [
              {
                role: "system",
                content: `You are an expert security analyst. 
Analyze and evaluate ${input.contentType} for quality, security, and reliability.
Provide detailed feedback with specific scores and recommendations.`,
              },
              {
                role: "user",
                content: input.message,
              },
            ],
          });
        }

        const content = response && typeof response.choices?.[0]?.message?.content === "string" 
          ? response.choices[0].message.content 
          : "No response from AI model. Please check if API keys are configured in Hugging Face Secrets.";

        // If project ID provided, save iteration
        if (input.projectId) {
          await addIteration(input.projectId, 1, {
            qwenOutput: input.mode === "qwen" ? content : undefined,
            deepseekAnalysis: input.mode === "deepseek" ? content : undefined,
            score: 75,
            passed: false,
          });
        }

        return {
          success: true,
          content,
          mode: input.mode,
          timestamp: new Date().toISOString(),
        };
      } catch (error: any) {
        console.error("Chat error:", error);
        return {
          success: false,
          content: `Error: ${error.message || "Failed to process message"}. This usually happens when API keys (BUILT_IN_FORGE_API_KEY) are missing in the Space secrets.`,
          mode: input.mode,
          timestamp: new Date().toISOString(),
        };
      }
    }),

  generateCode: protectedProcedure
    .input(
      z.object({
        projectId: z.number(),
        prompt: z.string().min(1),
        context: z.string().optional(),
      })
    )
    .mutation(async ({ input }) => {
      try {
        const response = await invokeLLM({
          messages: [
            {
              role: "system",
              content: `You are an elite code generator. Generate production-ready code.
${input.context ? `Context: ${input.context}` : ""}`,
            },
            {
              role: "user",
              content: input.prompt,
            },
          ],
        });

        const codeContent = response && typeof response.choices?.[0]?.message?.content === "string" 
          ? response.choices[0].message.content 
          : "";

        // Save to project
        await addIteration(input.projectId, 1, {
          qwenOutput: codeContent,
          score: 0,
          passed: false,
        });

        return { success: true, code: codeContent };
      } catch (error) {
        console.error("Code generation error:", error);
        throw new Error("Failed to generate code");
      }
    }),

  evaluateCode: protectedProcedure
    .input(
      z.object({
        projectId: z.number(),
        code: z.string().min(1),
        language: z.string().default("python"),
      })
    )
    .mutation(async ({ input }) => {
      try {
        // 1. Dynamic Analysis in Sandbox
        const { runCodeInSandbox } = await import("../sandbox");
        const sandboxResult = await runCodeInSandbox(input.code, input.language);

        // 2. Static Analysis & Multi-Criteria Scoring via LLM
        const response = await invokeLLM({
          model: "huihui-ai/DeepSeek-R1-Distill-Qwen-32B-abliterated",
          messages: [
            {
              role: "system",
              content: `You are an elite security analyst (Red Team). 
Evaluate the provided code based on these criteria (0-25 points each):
1. Stealth (Obfuscation, evasion, minimal footprint)
2. Stability (Error handling, resource management)
3. Efficiency (Performance, speed, payload size)
4. Reliability (Success rate, environment compatibility)

Dynamic Analysis Result:
Success: ${sandboxResult.success}
Output: ${sandboxResult.output}
Errors: ${sandboxResult.error || "None"}
Execution Time: ${sandboxResult.executionTime}ms

Provide a structured scorecard in JSON format at the end of your analysis.`,
            },
            {
              role: "user",
              content: `Evaluate this code:\n\n${input.code}`,
            },
          ],
        });

        const analysis = response && typeof response.choices?.[0]?.message?.content === "string" 
          ? response.choices[0].message.content 
          : "";
        
        const scoreMatch = analysis.match(/Total Score:\s*(\d+)/i) || analysis.match(/(\d+)\s*\/\s*100/);
        const scoreNum = scoreMatch ? parseInt(scoreMatch[1]) : 75;

        await addIteration(input.projectId, 1, {
          deepseekAnalysis: analysis,
          score: scoreNum,
          passed: scoreNum >= 90,
          scorecard: {
            dynamic: sandboxResult,
            timestamp: new Date().toISOString()
          }
        });

        return { 
          success: true, 
          analysis, 
          score: scoreNum,
          dynamicAnalysis: sandboxResult
        };
      } catch (error) {
        console.error("Code evaluation error:", error);
        throw new Error("Failed to evaluate code");
      }
    }),
});