@Service
public class CodeContextExtractor {

    private static final Logger log = LoggerFactory.getLogger(CodeContextExtractor.class);
    private static final int MAX_CONTEXT_CHARS = 8000; // 给上下文预留8000字符

    private final GitService gitService;
    private final ProjectStructureService structureService;

    public CodeContextExtractor(GitService gitService,
                                  ProjectStructureService structureService) {
        this.gitService = gitService;
        this.structureService = structureService;
    }

    /**
     * 为代码生成构建完整上下文
     */
    public CodeContext buildContext(CodeGenerationRequest request) {
        StringBuilder context = new StringBuilder();

        // 1. 项目基本信息
        ProjectInfo projectInfo = structureService.getProjectInfo(request.getProjectId());
        context.append("## 项目信息\n");
        context.append("框架：").append(projectInfo.getFramework()).append("\n");
        context.append("Java版本：").append(projectInfo.getJavaVersion()).append("\n");
        context.append("包名：").append(projectInfo.getBasePackage()).append("\n\n");

        // 2. 代码规范（从项目配置中读取）
        String codeStyle = projectInfo.getCodeStyle();
        if (codeStyle != null) {
            context.append("## 代码规范\n").append(codeStyle).append("\n\n");
        }

        // 3. 相关接口定义（如果是实现某个接口）
        if (request.getInterfaceName() != null) {
            String interfaceCode = gitService.getFileContent(
                    request.getProjectId(), request.getInterfaceName());
            if (interfaceCode != null) {
                context.append("## 需要实现的接口\n```java\n")
                        .append(truncate(interfaceCode, 2000))
                        .append("\n```\n\n");
            }
        }

        // 4. 当前文件中已有的代码（提供风格参考）
        if (request.getCurrentFileContent() != null) {
            context.append("## 当前文件已有代码\n```java\n")
                    .append(truncate(request.getCurrentFileContent(), 3000))
                    .append("\n```\n\n");
        }

        // 5. 相关依赖类的方法签名（不是完整实现，节省token）
        List<String> relatedClasses = request.getRelatedClasses();
        if (relatedClasses != null && !relatedClasses.isEmpty()) {
            context.append("## 相关类的方法签名\n");
            for (String className : relatedClasses) {
                String methodSignatures = extractMethodSignatures(
                        request.getProjectId(), className);
                if (methodSignatures != null) {
                    context.append("### ").append(className).append("\n")
                            .append(methodSignatures).append("\n\n");
                }
            }
        }

        return new CodeContext(context.toString(), request.getRequirement());
    }

    private String extractMethodSignatures(String projectId, String className) {
        String fullCode = gitService.getFileContent(projectId, className);
        if (fullCode == null) return null;

        // 提取方法签名（不含方法体）
        StringBuilder signatures = new StringBuilder();
        String[] lines = fullCode.split("\n");
        boolean inMethod = false;
        int braceDepth = 0;

        for (String line : lines) {
            String trimmed = line.trim();
            // 检测方法声明（简化版：public/private/protected开头且含括号）
            if (!inMethod && (trimmed.startsWith("public ") ||
                    trimmed.startsWith("private ") ||
                    trimmed.startsWith("protected ")) &&
                    trimmed.contains("(") && trimmed.endsWith("{")) {
                signatures.append(line.replace("{", ";")).append("\n");
                inMethod = true;
                braceDepth = 1;
            } else if (inMethod) {
                braceDepth += line.chars().filter(c -> c == '{').count();
                braceDepth -= line.chars().filter(c -> c == '}').count();
                if (braceDepth <= 0) {
                    inMethod = false;
                }
            }
        }
        return signatures.toString();
    }

    private String truncate(String text, int maxChars) {
        return text.length() <= maxChars ? text :
                text.substring(0, maxChars) + "\n... (已截断)";
    }

    @Data
    @AllArgsConstructor
    public static class CodeContext {
        private String contextText;
        private String requirement;
    }
}

3.2 代码生成服务

@Service
public class CodeGenerationService {

    private static final Logger log = LoggerFactory.getLogger(CodeGenerationService.class);

    private final ChatClient gpt4Client;
    private final CodeGeexClient codeGeexClient;
    private final CodeContextExtractor contextExtractor;
    private final CodePostProcessor postProcessor;
    private final CodeSecurityScanner securityScanner;

    private static final String CODE_GEN_SYSTEM_PROMPT = """
            你是一名资深Java工程师，专注于Spring Boot企业级应用开发。
            请根据需求生成高质量、可用于生产的Java代码。
            
            要求：
            1. 代码必须完整、可编译
            2. 包含必要的异常处理
            3. 添加JavaDoc注释
            4. 遵循项目已有的代码风格
            5. 只输出代码，用```java代码块包裹，不要有额外解释
            """;

    public CodeGenerationService(ChatClient.Builder builder,
                                   CodeGeexClient codeGeexClient,
                                   CodeContextExtractor contextExtractor,
                                   CodePostProcessor postProcessor,
                                   CodeSecurityScanner securityScanner) {
        this.gpt4Client = builder
                .defaultSystem(CODE_GEN_SYSTEM_PROMPT)
                .build();
        this.codeGeexClient = codeGeexClient;
        this.contextExtractor = contextExtractor;
        this.postProcessor = postProcessor;
        this.securityScanner = securityScanner;
    }

    /**
     * 生成完整的Java方法或类
     */
    public CodeGenerationResult generateCode(CodeGenerationRequest request) {
        // 1. 构建上下文
        CodeContextExtractor.CodeContext context =
                contextExtractor.buildContext(request);

        // 2. 路由到合适的模型
        String rawCode;
        String modelUsed;
        if (request.isComplexLogic()) {
            // 复杂业务逻辑用GPT-4o
            rawCode = generateWithGpt4(context);
            modelUsed = "gpt-4o";
        } else {
            // 简单补全用CodeGeeX（成本更低）
            rawCode = codeGeexClient.complete(buildCodeGeexPrompt(context));
            modelUsed = "codegeex";
        }

        // 3. 提取纯代码
        String extractedCode = postProcessor.extractCode(rawCode);

        // 4. 格式化
        String formattedCode = postProcessor.formatCode(extractedCode);

        // 5. 安全扫描
        List<SecurityIssue> securityIssues = securityScanner.scan(formattedCode);

        return CodeGenerationResult.builder()
                .code(formattedCode)
                .modelUsed(modelUsed)
                .securityIssues(securityIssues)
                .hasSecurityIssues(!securityIssues.isEmpty())
                .build();
    }

    private String generateWithGpt4(CodeContextExtractor.CodeContext context) {
        String userPrompt = """
                项目上下文：
                %s
                
                需求：
                %s
                
                请生成代码：
                """.formatted(context.getContextText(), context.getRequirement());

        return gpt4Client.prompt()
                .user(userPrompt)
                .call()
                .content();
    }

    private String buildCodeGeexPrompt(CodeContextExtractor.CodeContext context) {
        return context.getContextText() + "\n\n需求：" + context.getRequirement();
    }
}

3.3 代码后处理器

@Component
public class CodePostProcessor {

    private static final Pattern CODE_BLOCK_PATTERN =
            Pattern.compile("```(?:java|kotlin|groovy)?\\n([\\s\\S]*?)```",
                    Pattern.MULTILINE);

    /**
     * 从LLM输出中提取纯代码
     */
    public String extractCode(String llmOutput) {
        if (llmOutput == null || llmOutput.isEmpty()) return "";

        Matcher matcher = CODE_BLOCK_PATTERN.matcher(llmOutput);
        if (matcher.find()) {
            return matcher.group(1).trim();
        }

        // 没有代码块标记，尝试按启发式规则提取
        // 如果输出大部分都是代码（包含class/public等关键字），直接返回
        if (looksLikeCode(llmOutput)) {
            return llmOutput.trim();
        }

        // 找到第一个"class"或"public"开头的行，从那里截取
        String[] lines = llmOutput.split("\n");
        for (int i = 0; i < lines.length; i++) {
            if (lines[i].trim().startsWith("public class") ||
                    lines[i].trim().startsWith("@") ||
                    lines[i].trim().startsWith("package ")) {
                return String.join("\n",
                        Arrays.copyOfRange(lines, i, lines.length)).trim();
            }
        }

        return llmOutput.trim();
    }

    /**
     * 格式化Java代码（简单的缩进修正）
     */
    public String formatCode(String code) {
        if (code == null || code.isEmpty()) return code;
        // 实际项目可以集成Google Java Format或Palantir Java Format
        // 这里做简单的清理
        return code
                .replaceAll("\t", "    ")  // tab转4空格
                .replaceAll("\\r\\n", "\n")  // 统一换行符
                .replaceAll("[ \t]+\n", "\n")  // 删除行尾空白
                .trim();
    }

    private boolean looksLikeCode(String text) {
        String[] codeKeywords = {"public ", "private ", "class ", "@Override",
                "import ", "package ", "return ", "void "};
        int keywordCount = 0;
        for (String kw : codeKeywords) {
            if (text.contains(kw)) keywordCount++;
        }
        return keywordCount >= 3;
    }
}

3.4 代码安全扫描器

@Component
public class CodeSecurityScanner {

    // 常见安全问题的正则模式
    private static final List<SecurityRule> RULES = List.of(
            new SecurityRule("SQL_INJECTION",
                    Pattern.compile("executeQuery\\s*\\(\\s*\"[^\"]*\"\\s*\\+"),
                    "严重", "检测到字符串拼接SQL，可能存在SQL注入风险"),
            new SecurityRule("HARDCODED_PASSWORD",
                    Pattern.compile("(?i)password\\s*=\\s*\"[^\"]{4,}\""),
                    "严重", "检测到硬编码密码"),
            new SecurityRule("UNSAFE_DESERIALIZATION",
                    Pattern.compile("ObjectInputStream|readObject\\(\\)"),
                    "高危", "不安全的Java反序列化"),
            new SecurityRule("SYSTEM_EXIT",
                    Pattern.compile("System\\.exit\\s*\\("),
                    "中等", "使用了System.exit()，可能导致应用崩溃"),
            new SecurityRule("PATH_TRAVERSAL",
                    Pattern.compile("new File\\s*\\([^)]*\\+[^)]*\\)"),
                    "高危", "可能存在路径遍历漏洞"),
            new SecurityRule("WEAK_RANDOM",
                    Pattern.compile("new Random\\(\\)"),
                    "中等", "使用了不安全的Random，安全场景应使用SecureRandom")
    );

    public List<SecurityIssue> scan(String code) {
        if (code == null || code.isEmpty()) return List.of();

        List<SecurityIssue> issues = new ArrayList<>();
        for (SecurityRule rule : RULES) {
            Matcher matcher = rule.getPattern().matcher(code);
            while (matcher.find()) {
                int lineNumber = getLineNumber(code, matcher.start());
                issues.add(new SecurityIssue(
                        rule.getRuleId(),
                        rule.getSeverity(),
                        rule.getDescription(),
                        lineNumber,
                        code.substring(matcher.start(),
                                Math.min(matcher.end() + 20, code.length()))
                ));
            }
        }
        return issues;
    }

    private int getLineNumber(String code, int offset) {
        return (int) code.substring(0, offset).chars()
                .filter(c -> c == '\n').count() + 1;
    }

    @Data
    @AllArgsConstructor
    static class SecurityRule {
        private String ruleId;
        private Pattern pattern;
        private String severity;
        private String description;
    }
}

3.5 单元测试自动生成

@Service
public class TestGenerationService {

    private final ChatClient chatClient;

    private static final String TEST_GEN_PROMPT = """
            为以下Java方法生成完整的JUnit 5单元测试类。
            
            要求：
            1. 覆盖正常路径、边界条件、异常情况
            2. 使用Mockito Mock依赖
            3. 测试方法命名要清晰描述测试场景（如：shouldReturnUserWhenValidIdGiven）
            4. 每个测试方法只测一个场景
            5. 包含@BeforeEach初始化
            
            被测方法：
            ```java
            {method_code}
            ```
            
            相关依赖的接口：
            {dependencies}
            
            请生成完整测试类：
            """;

    public TestGenerationService(ChatClient.Builder builder) {
        this.chatClient = builder.build();
    }

    public String generateTests(String methodCode, String dependenciesInfo) {
        String prompt = TEST_GEN_PROMPT
                .replace("{method_code}", methodCode)
                .replace("{dependencies}", dependenciesInfo);

        String rawOutput = chatClient.prompt(prompt).call().content();
        return extractCode(rawOutput);
    }

    private String extractCode(String output) {
        Pattern pattern = Pattern.compile("```(?:java)?\\n([\\s\\S]*?)```");
        Matcher matcher = pattern.matcher(output);
        return matcher.find() ? matcher.group(1).trim() : output.trim();
    }
}

四、效果评估与优化

代码生成平台上线4个月，研发团队反馈数据：

指标	上线前	上线后
开发效率（功能点/人/天）	1.0（基准）	1.42
代码一次通过Code Review比例	68%	79%
单元测试覆盖率	61%	78%
AI生成代码可用率（无需大修）	-	71%
AI生成代码安全扫描通过率	-	94%

"AI生成代码可用率71%"表示71%的生成代码稍加修改就能直接使用，另外29%需要较大修改或完全重写。这个比例随着Prompt的持续优化在提升，第一个月只有53%，现在稳定在70%以上。

五、踩坑实录

坑1：代码上下文太长导致输出质量下降

一开始我把整个相关类都塞进上下文，以为信息越多越好。结果发现当上下文超过6000 token时，生成代码质量反而下降——模型似乎被大量细节"分心"了，生成的代码开始出现不相关的方法、错误的包名。后来把上下文限制在4000 token以内，只包含最相关的方法签名（不含方法体），生成质量明显回升。

坑2：代码提取正则漏掉了嵌套代码块

LLM有时候会在代码注释里再放一个代码示例，比如/** 示例：\n```\nfoo\n``` */。我的正则贪婪匹配，把这个注释里的代码块和后面真正的代码块全都提取了，拼接在一起，导致输出代码不完整甚至语法错误。改成非贪婪匹配（[\\s\\S]*?），同时在提取到多个代码块时，取最长的那个，这样通常是完整的类实现。

坑3：SecurityScanner的正则有大量误报

初版安全规则过于严格，比如ObjectInputStream的规则连项目里正当的反序列化（从受信任的内部数据源读取）都报警了。安全扫描报了满屏warning，开发者开始无视报警——破窗效应，真正的问题也被淹没了。后来把规则拆成了"阻断级"（直接拒绝合入）和"建议级"（提示审查），误报大幅减少，"阻断级"规则的信噪比达到了95%+。

六、总结

AI代码生成在Java企业开发中落地，关键不在于用哪个模型，而在于工程基础设施：上下文构建的质量、代码安全卫兵、持续的Prompt迭代。"一锤子买卖"的集成解决不了真实问题，要把代码生成当作一个持续演进的工程产品来做。