import httpx
import json
from httpx import BasicAuth
from time import sleep

class WorldQuantBrainAPI:
    def __init__(self, credentials_file='brain_credentials.txt'):
        self.credentials_file = credentials_file
        self.client = None
        self.brain_api_url = 'https://api.worldquantbrain.com'
    
    def load_credentials(self):
        """读取本地账号密码"""
        with open(self.credentials_file) as f:
            credentials = eval(f.read())
        return credentials[0], credentials[1]
    
    def login(self):
        """登录认证"""
        username, password = self.load_credentials()
        self.client = httpx.Client(auth=BasicAuth(username, password))
        
        response = self.client.post(f'{self.brain_api_url}/authentication')
        print(f"登录状态: {response.status_code}")
        
        if response.status_code == 201:
            print("登录成功!")
            return True
        else:
            print(f"登录失败: {response.json()}")
            return False
    
    def simulate_alpha(self, expression, settings=None):
        """模拟Alpha因子"""
        if self.client is None:
            raise Exception("请先登录")
        
        default_settings = {
            'instrumentType': 'EQUITY',
            'region': 'USA', 
            'universe': 'TOP3000',
            'delay': 1,
            'decay': 0,
            'neutralization': 'INDUSTRY',
            'truncation': 0.08,
            'pasteurization': 'ON',
            'unitHandling': 'VERIFY',
            'nanHandling': 'OFF',
            'language': 'FASTEXPR',
            'visualization': False,
        }
        
        if settings:
            default_settings.update(settings)
        
        simulation_data = {
            'type': 'REGULAR',
            'settings': default_settings,
            'regular': expression
        }
        
        sim_resp = self.client.post(
            f'{self.brain_api_url}/simulations',
            json=simulation_data,
        )
        print(f"模拟提交状态: {sim_resp.status_code}")
        
        sim_progress_url = sim_resp.headers['location']
        print(f"进度URL: {sim_progress_url}")
        
        while True:
            sim_progress_resp = self.client.get(sim_progress_url)
            retry_after_sec = float(sim_progress_resp.headers.get("Retry-After", 0))
            
            if retry_after_sec == 0:
                break
            print(sim_progress_resp.json())
            print(f"等待 {retry_after_sec} 秒...")
            sleep(retry_after_sec)
        
        alpha_id = sim_progress_resp.json()["alpha"]
        print(f"生成的Alpha ID: {alpha_id}")
        return alpha_id

    def close(self):
        """关闭连接"""
        if self.client:
            self.client.close()

if __name__ == "__main__":
    api = WorldQuantBrainAPI()
    
    try:
        # 读取账号密码并登录
        if api.login():
            # 模拟Alpha因子
            alpha_id = api.simulate_alpha("liabilities/assets")
            print(f"模拟完成，Alpha ID: {alpha_id}")
    finally:
        api.close()