#!/usr/bin/env python3
"""
千问AI代码评审系统 - 异步分批评审版本
支持大MR分多批处理，避免超时
"""

import os
import sys
import time

# 检查必要的依赖
try:
    import requests
    import json
    from datetime import datetime
except ImportError as e:
    print(f"❌ 缺少必要的Python库: {e}")
    print("请运行: pip3 install requests")
    sys.exit(1)


class QwenCodeReviewerAsync:
    def __init__(self):
        # 诊断环境
        self.mr_iid = os.getenv('CI_MERGE_REQUEST_IID')
        self.project_id = os.getenv('CI_PROJECT_ID')
        self.gitlab_url = os.getenv('CI_SERVER_URL', 'http://192.168.22.227')
        self.gitlab_token = os.getenv('AI_GITLAB_TOKEN')
        self.pipeline_id = os.getenv('CI_PIPELINE_ID', 'unknown')

        # 千问API配置
        self.api_url = "http://192.168.22.230:8000/v1/chat/completions"
        self.model_name = "Qwen3-32B-AWQ"

        # 分批配置
        self.batch_size = int(os.getenv('REVIEW_BATCH_SIZE', '3'))  # 每批文件数
        self.max_total_files = int(os.getenv('MAX_REVIEW_FILES', '50'))  # 最多评审文件数
        self.max_diff_per_file = int(os.getenv('MAX_DIFF_PER_FILE', '8000'))  # 单文件diff限制

        print(f"🔍 环境检测: MR#{self.mr_iid}, 项目#{self.project_id}")
        print(f"⚙️  分批配置: 每批{self.batch_size}个文件, 最多{self.max_total_files}个文件")

    def get_mr_changes(self):
        """获取MR的真实代码差异"""
        if not all([self.mr_iid, self.project_id, self.gitlab_token]):
            print("❌ 缺少必要的MR环境变量或Token")
            return None

        print(f"📥 获取MR #{self.mr_iid} 的代码变更...")
        url = f"{self.gitlab_url}/api/v4/projects/{self.project_id}/merge_requests/{self.mr_iid}/changes"
        headers = {"PRIVATE-TOKEN": self.gitlab_token}

        try:
            response = requests.get(url, headers=headers, timeout=30)
            response.raise_for_status()
            data = response.json()
            changes = data.get('changes', [])

            if not changes:
                print("⚠️  MR中没有代码变更")
                return []

            print(f"✅ 获取到 {len(changes)} 个文件变更")
            return changes
        except Exception as e:
            print(f"❌ 获取MR变更失败: {e}")
            return None

    def prioritize_files(self, changes):
        """按优先级排序文件，核心代码优先，跳过不需要评审的文件"""
        priority_map = {
            '.c': 1, '.h': 1, '.cpp': 1,  # 核心代码最高优先级
            '.py': 2, '.sh': 2, '.js': 2,  # 脚本次之
            '.json': 3,  # 配置
        }

        # 跳过不需要评审的文件类型
        skip_exts = {'.yml', '.yaml', '.md', '.txt', '.rst', '.log'}

        def get_priority(change):
            path = change.get('new_path') or change.get('old_path', '')
            ext = os.path.splitext(path)[1].lower()
            if ext in skip_exts:
                return 99  # 跳过的文件给最低优先级
            return priority_map.get(ext, 2)  # 默认优先级2

        # 过滤并排序（排除优先级99的文件）
        sorted_changes = [c for c in sorted(changes, key=get_priority) if get_priority(c) != 99]

        # 限制总数
        limited_changes = sorted_changes[:self.max_total_files]

        # 打印文件列表
        print(f"\n📋 文件优先级排序（前{len(limited_changes)}个）:")
        for i, change in enumerate(limited_changes, 1):
            path = change.get('new_path') or change.get('old_path', '未知')
            priority = get_priority(change)
            priority_label = {1: '🔴核心', 2: '🟡脚本', 3: '🔵配置'}.get(priority, '?')
            print(f"   {i:2d}. {priority_label} {path}")

        skipped = len(changes) - len(limited_changes)
        if skipped > 0:
            print(f"   ... 还有 {skipped} 个文件未列入评审（超出限制）")

        return limited_changes

    def create_batch_prompt(self, batch_changes, batch_num, total_batches, project_context):
        """为一批文件创建评审提示词"""

        # 精简的系统提示词（缩短长度）
        system_prompt = """你是资深代码架构师，按以下优先级评审代码：

🔴 P0-阻塞：功能缺陷、安全漏洞、资源泄漏、并发问题
🟠 P1-重要：性能瓶颈、错误处理、圈复杂度过高、重复代码  
🟡 P2-建议：命名规范、缺少注释、硬编码
⚪ P3-可选：格式问题、魔法数字

输出格式：
## 优先级 问题简述
**位置**: `文件#L行`
**原因**: 一句话说明
**修复建议**: `代码示例`

要求：
- 每个问题必给具体代码和修复方案
- 优先报告P0/P1问题
- 每个优先级最多3个问题"""

        # 构建提示词
        prompt = f"""{system_prompt}

---
项目上下文：
- 分支: {project_context.get('branch', 'Unknown')}
- 作者: {project_context.get('author', 'Unknown')}
- 批次: 第{batch_num}批/共{total_batches}批
- 本批文件数: {len(batch_changes)}
---

待评审代码变更：
"""

        total_diff_len = 0
        for i, change in enumerate(batch_changes, 1):
            file_path = change.get('new_path') or change.get('old_path', '未知文件')
            original_diff = change.get('diff', '')

            # 截断单文件diff
            if len(original_diff) > self.max_diff_per_file:
                diff = original_diff[:self.max_diff_per_file]
                diff += f"\n... (截断，原长度{len(original_diff)}) ..."
            else:
                diff = original_diff

            total_diff_len += len(diff)

            prompt += f"""
## 文件{i}: `{file_path}`
```diff
{diff}
```
"""

        prompt += "\n请按上述标准进行专业评审："

        print(f"\n📝 第{batch_num}批提示词: {len(prompt)} 字符 (diff内容: {total_diff_len} 字符)")
        return prompt

    def call_qwen_api(self, prompt, max_retries=3, base_timeout=90):
        """调用千问API，带重试机制"""
        headers = {"Content-Type": "application/json"}
        data = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": "你是代码评审专家，输出结构化评审报告。"},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.2,
            "max_tokens": 3000  # 减少输出token，加快响应
        }

        for attempt in range(1, max_retries + 1):
            try:
                print(f"   🔄 API调用尝试 {attempt}/{max_retries}...")
                timeout = base_timeout + (attempt - 1) * 30
                response = requests.post(self.api_url, headers=headers, json=data, timeout=timeout)
                response.raise_for_status()
                result = response.json()
                content = result['choices'][0]['message']['content']
                if attempt > 1:
                    print(f"   ✅ 第 {attempt} 次尝试成功！")
                return content
            except requests.exceptions.Timeout:
                print(f"   ⏱️  第 {attempt} 次超时 (timeout={timeout}s)")
                if attempt < max_retries:
                    wait_time = min(2 ** attempt, 10)
                    print(f"   ⏳ 等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                else:
                    return f"❌ 本批评审超时，请稍后重试或联系管理员"
            except Exception as e:
                return f"❌ API调用失败: {str(e)}"

    def post_pending_notice(self):
        """发布"评审中"通知"""
        if not all([self.mr_iid, self.project_id, self.gitlab_token]):
            return None

        url = f"{self.gitlab_url}/api/v4/projects/{self.project_id}/merge_requests/{self.mr_iid}/notes"
        headers = {
            "PRIVATE-TOKEN": self.gitlab_token,
            "Content-Type": "application/json"
        }

        notice = f"""## 🤖 千问AI代码评审进行中

⏳ **状态**: 正在分析代码变更...
📊 **配置**: 每批 {self.batch_size} 个文件，最多评审 {self.max_total_files} 个文件
⏱️ **预计**: 2-5 分钟完成（根据文件数量）

> 本评审由 AI 自动生成，请勿重复触发。
> 流水线: {self.pipeline_id}"""

        try:
            response = requests.post(url, headers=headers, json={"body": notice}, timeout=30)
            response.raise_for_status()
            note_id = response.json().get('id')
            print(f"✅ 已发布评审中通知 (note_id: {note_id})")
            return note_id
        except Exception as e:
            print(f"⚠️  发布通知失败: {e}")
            return None

    def update_review_result(self, note_id, review_content, batch_info):
        """更新评审结果（编辑原评论）"""
        if not all([self.mr_iid, self.project_id, self.gitlab_token, note_id]):
            # 如果没有note_id，创建新评论
            return self.post_new_review(review_content, batch_info)

        url = f"{self.gitlab_url}/api/v4/projects/{self.project_id}/merge_requests/{self.mr_iid}/notes/{note_id}"
        headers = {
            "PRIVATE-TOKEN": self.gitlab_token,
            "Content-Type": "application/json"
        }

        full_content = f"""## 🤖 千问AI代码评审报告 (Qwen3-32B-AWQ)

{batch_info}

{review_content}

---
**评审信息**
- **MR**: !{self.mr_iid}
- **模型**: {self.model_name}
- **完成时间**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
- **触发**: GitLab CI/CD 流水线

> 💡 本评审由公司内部千问大模型自动生成，旨在辅助代码质量提升。请结合人工评审做出最终决策。"""

        try:
            response = requests.put(url, headers=headers, json={"body": full_content}, timeout=30)
            response.raise_for_status()
            print("✅ 已更新评审结果")
            return True
        except Exception as e:
            print(f"⚠️  更新评论失败: {e}，尝试创建新评论")
            return self.post_new_review(review_content, batch_info)

    def post_new_review(self, review_content, batch_info):
        """创建新的评审评论"""
        url = f"{self.gitlab_url}/api/v4/projects/{self.project_id}/merge_requests/{self.mr_iid}/notes"
        headers = {
            "PRIVATE-TOKEN": self.gitlab_token,
            "Content-Type": "application/json"
        }

        full_content = f"""## 🤖 千问AI代码评审报告 (Qwen3-32B-AWQ)

{batch_info}

{review_content}

---
**评审信息**
- **MR**: !{self.mr_iid}
- **模型**: {self.model_name}
- **完成时间**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
- **触发**: GitLab CI/CD 流水线

> 💡 本评审由公司内部千问大模型自动生成，旨在辅助代码质量提升。"""

        try:
            response = requests.post(url, headers=headers, json={"body": full_content}, timeout=30)
            response.raise_for_status()
            print("✅ 已创建新评审评论")
            return True
        except Exception as e:
            print(f"❌ 发布评审失败: {e}")
            return False

    def run(self):
        """主执行流程 - 异步分批评审"""
        print("=" * 60)
        print("🚀 启动千问AI代码评审（异步分批模式）")
        print("=" * 60)

        # 1. 获取MR变更
        changes = self.get_mr_changes()
        if changes is None:
            print("❌ 无法获取MR变更，退出")
            return False
        if not changes:
            print("⚠️  MR中没有代码变更")
            return True

        # 2. 文件优先级排序
        prioritized_changes = self.prioritize_files(changes)

        # 3. 先发布"评审中"通知
        note_id = self.post_pending_notice()

        # 4. 分批处理
        total_files = len(prioritized_changes)
        batch_size = self.batch_size
        total_batches = (total_files + batch_size - 1) // batch_size

        all_reviews = []
        project_context = {
            'branch': os.getenv('CI_COMMIT_REF_NAME', 'Unknown'),
            'author': os.getenv('GITLAB_USER_NAME', 'Unknown'),
        }

        print(f"\n{'=' * 60}")
        print(f"📦 开始分批评审: 共 {total_files} 个文件，分 {total_batches} 批")
        print(f"{'=' * 60}")

        for batch_num in range(1, total_batches + 1):
            start_idx = (batch_num - 1) * batch_size
            end_idx = min(start_idx + batch_size, total_files)
            batch_changes = prioritized_changes[start_idx:end_idx]

            # 显示批次信息：共X批，当前第Y批，文件a到b
            print(f"\n📌 批次进度: 共{total_batches}批 | 当前第{batch_num}批 | 文件{start_idx+1}-{end_idx}/{total_files}")
            print("-" * 60)

            # 显示本批文件列表
            for i, change in enumerate(batch_changes, start_idx + 1):
                file_path = change.get('new_path') or change.get('old_path', '未知')
                print(f"   [{i:2d}] {os.path.basename(file_path)}")

            # 创建提示词
            prompt = self.create_batch_prompt(batch_changes, batch_num, total_batches, project_context)

            # 调用AI
            review_result = self.call_qwen_api(prompt)

            # 保存结果
            all_reviews.append({
                'batch_num': batch_num,
                'files': [c.get('new_path') or c.get('old_path') for c in batch_changes],
                'result': review_result
            })

            # 批次间短暂休息，避免压垮模型
            if batch_num < total_batches:
                print("   😴 批次间休息 2 秒...")
                time.sleep(2)

        # 5. 合并所有评审结果
        print(f"\n{'=' * 60}")
        print("📊 合并所有批次评审结果...")
        print(f"{'=' * 60}")

        merged_review = self.merge_reviews(all_reviews, total_files, len(changes))

        # 6. 发布最终结果
        batch_info = f"📦 评审范围: 共 {len(changes)} 个文件变更，详细评审了 {total_files} 个文件（分 {total_batches} 批）"
        success = self.update_review_result(note_id, merged_review, batch_info)

        if success:
            print("\n🎉 AI代码评审流程完成！")
        else:
            print("\n⚠️  评审完成但发布失败，请查看日志")

        return success

    def merge_reviews(self, all_reviews, reviewed_count, total_count):
        """合并多批评审结果"""
        merged = []

        # 批次详情
        merged.append("### 📋 分批评审详情\n")
        for review in all_reviews:
            file_list = ', '.join([os.path.basename(f) for f in review['files'][:3]])
            if len(review['files']) > 3:
                file_list += f" 等{len(review['files'])}个文件"
            merged.append(f"- **第{review['batch_num']}批**: {file_list}")
        merged.append("")

        # 合并各批次内容
        merged.append("---\n")
        for review in all_reviews:
            merged.append(f"\n### 第 {review['batch_num']} 批评审结果\n")
            merged.append(review['result'])
            merged.append("\n")

        # 统计信息
        skipped = total_count - reviewed_count
        if skipped > 0:
            merged.append(f"\n> ⚠️  因长度限制，{skipped} 个文件未详细评审")

        return '\n'.join(merged)


if __name__ == "__main__":
    reviewer = QwenCodeReviewerAsync()
    success = reviewer.run()
    sys.exit(0 if success else 1)