# -*- coding: utf-8 -*-
import os
import httpx
from requests.auth import HTTPBasicAuth
from urllib.parse import urlencode


class BrainLogin:
    def __init__(self, credentials_file='account.txt'):
        self.credentials_file = credentials_file
        self.client = None
        self.brain_api_url = 'https://api.worldquantbrain.com'

    def load_credentials(self):
        if not os.path.exists(self.credentials_file):
            print("未找到 account.txt 文件")
            with open(self.credentials_file, 'w') as f:
                f.write("")
            print("account.txt 文件已创建，请填写账号密码, 格式: ['username', 'password']")
            exit(1)

        with open(self.credentials_file) as f:
            credentials = eval(f.read())
        return credentials[0], credentials[1]

    def login(self):
        try:
            username, password = self.load_credentials()
            self.client = httpx.Client(auth=HTTPBasicAuth(username, password))

            response = self.client.post(f'{self.brain_api_url}/authentication')
            print(f"登录状态: {response.status_code}")

            if response.status_code in [200, 201]:
                print("登录成功!")
                return self.client
            else:
                print(f"登录失败: {response.json()}")
                return None

        except Exception as e:
            print(f"登录过程中出现错误: {e}")
            return None

    def get_alphas(self, **params):
        if self.client is None:
            print("请先登录!")
            return None

        try:
            encoded_params = urlencode(params, doseq=True)
            url = f"{self.brain_api_url}/users/self/alphas?{encoded_params}"

            response = self.client.get(url)

            if response.status_code == 200:
                return response.json()
            else:
                print(f"获取alpha列表失败: {response.text}")
                return None

        except Exception as e:
            print(f"获取alpha列表过程中出现错误: {e}")
            return None


# 使用示例
if __name__ == "__main__":
    brain = BrainLogin()

    if brain.login():
        alphas = brain.get_alphas(
            limit=50,
            offset=0,
            status="UNSUBMITTED",
            order="dateSubmitted",
            hidden="false"
        )

        if alphas:
            print("获取alpha列表成功!")
            print(alphas)
            print(f"一共 {len(alphas)} 个 alpha")