import json
import time
from typing import Dict, Any, Optional

import requests


class APIClient:
    """DeepSeek API 客户端，含重试逻辑。"""

    def __init__(self, api_key: str, model: str = 'deepseek-v4-flash',
                 base_url: str = 'https://api.deepseek.com/chat/completions',
                 max_retries: int = 3, timeout: int = 120):
        self.api_key = api_key
        self.model = model
        self.base_url = base_url
        self.max_retries = max_retries
        self.timeout = timeout

    def generate(self, prompt: str) -> Dict[str, Any]:
        """发送 prompt 并返回 AI 生成的结果。"""
        system_prompt = (
            "你是COBOL程序的测试数据生成专家。"
            "请严格按照提供的规则，生成符合格式要求的测试数据。"
            "输出必须是可被json.loads()直接解析的JSON，不要包裹在```json```代码块中。"
            "不要在JSON前后添加任何说明文字。"
        )

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ]

        last_error = None

        for attempt in range(1, self.max_retries + 1):
            try:
                response = self._call_api(messages)
                content = response['choices'][0]['message']['content']
                data = self._parse_json(content)
                if data is not None:
                    return data

                error_msg = (
                    f"前回の出力は有効なJSONではありませんでした。"
                    f"必ず有効なJSONのみを出力してください。"
                    f"コードブロック（```）で囲まないでください。"
                )
                messages.append({"role": "assistant", "content": content})
                messages.append({"role": "user", "content": error_msg})

            except requests.exceptions.RequestException as e:
                last_error = e
                if attempt < self.max_retries:
                    time.sleep(2 ** attempt)
                continue

        raise RuntimeError(
            f"API调用失败，已重试{self.max_retries}次。"
            f"最后错误: {last_error}"
        )

    def _call_api(self, messages: list) -> dict:
        """单次 API 调用。"""
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json',
        }
        payload = {
            'model': self.model,
            'messages': messages,
            'temperature': 0.3,
            'max_tokens': 8192,
        }

        resp = requests.post(
            self.base_url,
            headers=headers,
            json=payload,
            timeout=self.timeout
        )
        resp.raise_for_status()
        return resp.json()

    @staticmethod
    def _parse_json(text: str) -> Optional[dict]:
        """尝试从文本中提取 JSON。"""
        text = text.strip()

        if text.startswith('```json'):
            text = text[7:]
        if text.startswith('```'):
            text = text[3:]
        if text.endswith('```'):
            text = text[:-3]
        text = text.strip()

        try:
            return json.loads(text)
        except json.JSONDecodeError:
            start = text.find('{')
            end = text.rfind('}')
            if start >= 0 and end > start:
                try:
                    return json.loads(text[start:end + 1])
                except json.JSONDecodeError:
                    pass
            return None