from django.core.management.base import BaseCommand
from openai import OpenAI
from huodong.models import Branch
import time
import re


class Command(BaseCommand):
    help = 'Update branch provinces using AI API'

    def handle(self, *args, **options):
        client = OpenAI(
            base_url="https://integrate.api.nvidia.com/v1",
            api_key="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"
        )

        def get_correct_province(branch_name, current_location, retry_count=0):
            prompt = f"""请根据以下分支机构名称，返回一个正确的中国省份名称（全称，如：浙江省、北京市、上海市、广东省等）。

分支机构名称：{branch_name}

要求：
1. 只返回省份名称，不要包含任何其他文字
2. 省份名称必须是中国的省级行政区全称
3. 如果无法确定，返回None
例子一：
分支机构名称：石家庄中山西路营业部
返回：河北省
例子二：
分支机构名称：北京市海淀区营业部
返回：北京市
例子三：
分支机构名称：奎屯北京东路营业部
返回：新疆维吾尔自治区

"""


            try:
                completion = client.chat.completions.create(
                    model="deepseek-ai/deepseek-r1",
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.3,
                    top_p=0.7,
                    max_tokens=100,
                    stream=False
                )

                message = completion.choices[0].message
                reasoning = getattr(message, "reasoning_content", None)
                content = message.content

                if reasoning:
                    self.stdout.write(f"  推理过程: {reasoning[:200]}...")
                
                if content is None:
                    if reasoning:
                        self.stdout.write(f"  content为None，尝试从推理过程提取...")
                        
                        patterns = [
                            r'(?:正确的(?:答案)?应该是|所以|因此|答案是|结果为)[:：\s]*(北京市|天津市|上海市|重庆市|河北省|山西省|辽宁省|吉林省|黑龙江省|江苏省|浙江省|安徽省|福建省|江西省|山东省|河南省|湖北省|湖南省|广东省|海南省|四川省|贵州省|云南省|陕西省|甘肃省|青海省|台湾省|内蒙古自治区|广西壮族自治区|西藏自治区|宁夏回族自治区|新疆维吾尔自治区|香港特别行政区|澳门特别行政区)',
                            r'(北京市|天津市|上海市|重庆市|河北省|山西省|辽宁省|吉林省|黑龙江省|江苏省|浙江省|安徽省|福建省|江西省|山东省|河南省|湖北省|湖南省|广东省|海南省|四川省|贵州省|云南省|陕西省|甘肃省|青海省|台湾省|内蒙古自治区|广西壮族自治区|西藏自治区|宁夏回族自治区|新疆维吾尔自治区|香港特别行政区|澳门特别行政区)'
                        ]
                        
                        for pattern in patterns:
                            match = re.search(pattern, reasoning)
                            if match:
                                result = match.group(1)
                                self.stdout.write(f"  从推理提取结果: {result}")
                                return result
                    
                    if retry_count < 2:
                        self.stdout.write(f"  content为None，重试中 ({retry_count + 1}/2)...")
                        time.sleep(2)
                        return get_correct_province(branch_name, current_location, retry_count + 1)
                    else:
                        self.stdout.write(self.style.WARNING(f"  多次重试后仍无结果，使用当前省份"))
                        return current_location
                
                result = content.strip()
                self.stdout.write(f"  原始结果: {result}")
                
                return result
                
            except Exception as e:
                if retry_count < 2:
                    self.stdout.write(self.style.WARNING(f"  API调用失败: {e}，重试中 ({retry_count + 1}/2)..."))
                    time.sleep(2)
                    return get_correct_province(branch_name, current_location, retry_count + 1)
                else:
                    self.stdout.write(self.style.ERROR(f"  多次重试后仍失败: {e}，使用当前省份"))
                    return current_location

        branches = Branch.objects.all()
        total = branches.count()
        updated_count = 0
        
        self.stdout.write(self.style.SUCCESS(f"开始处理 {total} 个分支机构..."))
        self.stdout.write("=" * 80)
        
        for index, branch in enumerate(branches, 1):
            self.stdout.write(f"\n[{index}/{total}] 处理: {branch.name}")
            self.stdout.write(f"当前省份: {branch.location}")
            
            try:
                correct_province = get_correct_province(branch.name, branch.location)
                self.stdout.write(f"建议省份: {correct_province}")
                
                if correct_province != branch.location:
                    old_location = branch.location
                    branch.location = correct_province
                    branch.save()
                    self.stdout.write(self.style.SUCCESS(f"✓ 已更新: {old_location} -> {correct_province}"))
                    updated_count += 1
                else:
                    self.stdout.write("- 省份未变化")
                    
            except Exception as e:
                self.stdout.write(self.style.ERROR(f"✗ 处理失败: {e}"))
            
            self.stdout.write("-" * 80)
        
        self.stdout.write(self.style.SUCCESS(f"\n处理完成！"))
        self.stdout.write(f"总计: {total} 个分支机构")
        self.stdout.write(f"已更新: {updated_count} 个")
        self.stdout.write(f"未变化: {total - updated_count} 个")