#!/usr/bin/env python3
"""
优化的DeepSeek v3.1代码规约化测试

减少超时，优化测试流程
"""

import os
import requests
import json
import time


class OptimizedDeepSeekTest:
    """优化的DeepSeek测试类"""

    def __init__(self, api_key=None, base_url=None):
        self.api_key = api_key or os.getenv('SILICONFLOW_API_KEY')
        self.base_url = base_url or 'https://api.siliconflow.cn/v1'
        self.model = 'deepseek-ai/DeepSeek-V3.1'

        if not self.api_key:
            raise ValueError("需要设置 SILICONFLOW_API_KEY 环境变量")

    def generate_specification(self, function_info, verification_goals):
        """生成形式化规约"""
        prompt = self._build_prompt(function_info, verification_goals)

        messages = [
            {
                "role": "system",
                "content": "你是CBMC形式化验证专家，为C/C++代码生成准确的CBMC规约。只输出CBMC代码，不要解释。"
            },
            {
                "role": "user",
                "content": prompt
            }
        ]

        data = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.1,
            "max_tokens": 1024,
            "stream": False
        }

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        start_time = time.time()

        try:
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=data,
                timeout=60
            )

            if response.status_code == 200:
                result = response.json()
                content = result['choices'][0]['message']['content']
                tokens_used = result.get('usage', {}).get('total_tokens', 0)

                # 简单的后处理
                specification = self._extract_cbmc_spec(content)

                generation_time = time.time() - start_time

                return {
                    'specification': specification,
                    'raw_content': content,
                    'generation_time': generation_time,
                    'tokens_used': tokens_used,
                    'success': True
                }
            else:
                return {
                    'success': False,
                    'error': f"API错误: {response.status_code}"
                }

        except Exception as e:
            return {
                'success': False,
                'error': f"请求失败: {str(e)}"
            }

    def _build_prompt(self, function_info, verification_goals):
        """构建简洁的提示词"""
        func_name = function_info.get('name', 'unknown')
        return_type = function_info.get('return_type', 'void')
        parameters = function_info.get('parameters', [])
        source_code = function_info.get('source_code', '')

        # 构建函数签名
        param_str = ', '.join([f"{p['type']} {p['name']}" for p in parameters])
        function_signature = f"{return_type} {func_name}({param_str})"

        prompt = f"""为以下C函数生成CBMC规约：

{function_signature}

源代码：
{source_code}

验证目标：{', '.join(verification_goals)}

只输出CBMC规约代码：
\\requires ...
\\ensures ...
\\assigns ..."""
        return prompt

    def _extract_cbmc_spec(self, content):
        """提取CBMC规约"""
        lines = content.split('\n')
        cbmc_lines = []

        for line in lines:
            line = line.strip()
            if (line.startswith(r'\requires') or
                line.startswith(r'\ensures') or
                line.startswith(r'\assigns') or
                line.startswith(r'\valid')):
                cbmc_lines.append(line)

        return '\n'.join(cbmc_lines) if cbmc_lines else content.strip()


def create_test_functions():
    """创建测试函数"""
    return [
        {
            'name': 'linked_list_insert',
            'return_type': 'bool',
            'parameters': [
                {'name': 'head', 'type': 'struct Node**'},
                {'name': 'value', 'type': 'int'}
            ],
            'source_code': '''bool linked_list_insert(struct Node** head, int value) {
    if (head == NULL) return false;

    struct Node* new_node = (struct Node*)malloc(sizeof(struct Node));
    if (new_node == NULL) return false;

    new_node->data = value;
    new_node->next = *head;
    *head = new_node;
    return true;
}''',
            'verification_goals': ['functional_correctness', 'memory_safety', 'pointer_validity', 'error_handling']
        },
        {
            'name': 'string_copy_safe',
            'return_type': 'size_t',
            'parameters': [
                {'name': 'dest', 'type': 'char*'},
                {'name': 'src', 'type': 'const char*'},
                {'name': 'dest_size', 'type': 'size_t'}
            ],
            'source_code': '''size_t string_copy_safe(char* dest, const char* src, size_t dest_size) {
    if (dest == NULL || src == NULL || dest_size == 0) {
        return 0;
    }

    size_t i;
    for (i = 0; i < dest_size - 1 && src[i] != '\\0'; i++) {
        dest[i] = src[i];
    }
    dest[i] = '\\0';
    return i;
}''',
            'verification_goals': ['functional_correctness', 'memory_safety', 'array_bounds', 'error_handling']
        }
    ]


def main():
    """主测试函数"""
    print("🚀 DeepSeek v3.1 优化代码规约化测试")
    print("=" * 50)

    try:
        # 创建测试实例
        tester = OptimizedDeepSeekTest()
        print(f"📡 模型: {tester.model}")

        # 测试函数
        test_functions = create_test_functions()

        for i, func in enumerate(test_functions, 1):
            print(f"\n{'='*20} 测试 {i} {'='*20}")
            print(f"📝 函数: {func['name']}")
            print(f"🎯 目标: {', '.join(func['verification_goals'])}")

            print("\n⏳ 生成中...")
            result = tester.generate_specification(func, func['verification_goals'])

            if result['success']:
                print(f"✅ 成功! ⏱️ {result['generation_time']:.2f}s, 🔤 {result['tokens_used']} tokens")

                print(f"\n📋 CBMC规约:")
                print("-" * 40)
                print(result['specification'])
                print("-" * 40)

                print(f"\n📄 原始输出:")
                print("-" * 30)
                print(result['raw_content'])
                print("-" * 30)
            else:
                print(f"❌ 失败: {result['error']}")

        print(f"\n🎉 优化测试完成!")

    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")


if __name__ == "__main__":
    main()