import time
import json
import pandas as pd
import numpy as np
from openai import OpenAI
from .schemas import DatasetSchema

class SDGEngine:
    def __init__(self, api_key: str):
        if not api_key:
            raise ValueError("OpenAI API Key is required.")
        self.client = OpenAI(api_key=api_key)
        self.schema = None
        self.stats = {}

    def fit(self, file_path: str):
        """分析原始数据以提取约束和业务逻辑。"""
        df = pd.read_csv(file_path)
        for col in df.select_dtypes(include=[np.number]).columns:
            self.stats[col] = {
                "min": float(df[col].min()), 
                "max": float(df[col].max())
            }
        
        print("🧠 [1/5] 正在提取数据逻辑和数值边界...")
        sample = df.head(5).to_markdown()
        completion = self.client.beta.chat.completions.parse(
            model="gpt-4o-2024-08-06",
            messages=[{"role": "user", "content": f"Analyze this dataset structure:\n{sample}"}],
            response_format=DatasetSchema,
        )
        self.schema = completion.choices[0].message.parsed
        print(f"✨ 逻辑提取完成: {self.schema.description[:60]}...")

    def auto_generate(self, source_csv: str, count: int):
        """完整流水线：分析 -> 云端批处理 -> 数据重构。"""
        self.fit(source_csv)
        
        batch_file = "batch_tasks.jsonl"
        print(f"📦 [2/5] 正在构建包含逻辑约束的 {count} 条任务...")
        constraints = ", ".join([f"{k} must be between {v['min']} and {v['max']}" for k, v in self.stats.items()])
        
        with open(batch_file, "w", encoding="utf-8") as f:
            for i in range(count):
                task = {
                    "custom_id": f"row-{i}", 
                    "method": "POST", 
                    "url": "/v1/chat/completions",
                    "body": {
                        "model": "gpt-4o-mini",
                        "messages": [
                            {"role": "system", "content": f"Expert SDG. Context: {self.schema.description}. Constraints: {constraints}. Do not generate PII."},
                            {"role": "user", "content": "Generate 1 synthetic JSON record."}
                        ],
                        "response_format": {"type": "json_object"}
                    }
                }
                f.write(json.dumps(task) + "\n")

        print("🚀 [3/5] 通过 Batch API 启动云端并行合成...")
        cloud_file = self.client.files.create(file=open(batch_file, "rb"), purpose="batch")
        job = self.client.batches.create(
            input_file_id=cloud_file.id, 
            endpoint="/v1/chat/completions", 
            completion_window="24h"
        )

        print("⏳ [4/5] 等待云端结果（异步模式）...")
        while True:
            status = self.client.batches.retrieve(job.id)
            if status.status == "completed":
                print("\n✅ 云端合成成功完成！")
                break
            elif status.status in ["failed", "expired", "cancelled"]:
                raise Exception(f"批处理任务失败，状态为: {status.status}")
            
            progress = status.request_counts
            print(f"   进度: 已完成 {progress.completed}/{progress.total} 条记录...", end="\r")
            time.sleep(10)

        print("📥 [5/5] 正在下载并重构 Pandas DataFrame...")
        result_text = self.client.files.content(status.output_file_id).text
        raw_rows = []
        for line in result_text.strip().split('\n'):
            res = json.loads(line)
            content = json.loads(res['response']['body']['choices'][0]['message']['content'])
            row = list(content.values())[0] if isinstance(content, dict) and len(content) == 1 else content
            raw_rows.append(row)
        
        return pd.DataFrame(raw_rows)
