# -*- coding: utf-8 -*-
import random
import time
from xmlrpc.client import ServerProxy, Fault

is_local = 0
if is_local:
    url = 'http://127.0.0.1:28888'
else:
    url = 'http://192.168.31.41:32000'

username = 'rpc'
password = 'aaaAAA111'
db = "quantify"

common = ServerProxy(f"{url}/xmlrpc/2/common")
uid = common.authenticate(db, username, password, {})

models = ServerProxy(f"{url}/xmlrpc/2/object")

# 1. 在 alpha.generate.research.direction 表中创建一条数据
# 需要准备: region_id 和 universe_id

region = 'USA'
universe = 'TOP3000'

# 根据名称搜索 region 和 universe 的ID
region_ids = models.execute_kw(db, uid, password, 'alpha.region.settings', 'search', [[['name', '=', region]]], {'limit': 1})
universe_ids = models.execute_kw(db, uid, password, 'alpha.universe.settings', 'search', [[['name', '=', universe]]], {'limit': 1})

if not region_ids:
    print(f"没有找到 region: '{region}'，请检查名称是否正确")
    exit(1)
if not universe_ids:
    print(f"没有找到 universe: '{universe}'，请检查名称是否正确")
    exit(1)

region_id = region_ids[0]
universe_id = universe_ids[0]

# 获取 region 和 universe 的 name
region_data = models.execute_kw(db, uid, password, 'alpha.region.settings', 'read', [[region_id], ['name']])
universe_data = models.execute_kw(db, uid, password, 'alpha.universe.settings', 'read', [[universe_id], ['name']])
region_name = region_data[0]['name'] if region_data else 'Unknown'
universe_name = universe_data[0]['name'] if universe_data else 'Unknown'

print(f"使用 region_id: {region_id} (name: {region_name}), universe_id: {universe_id} (name: {universe_name})")

# 随机选择一个 llm_settings_line_id
llm_line_ids = models.execute_kw(db, uid, password, 'llm.settings.line', 'search', [[['id', '!=', 0]]])
if not llm_line_ids:
    print("没有找到 llm.settings.line 记录")
    exit(1)

llm_settings_line_id = random.choice(llm_line_ids)
llm_line_data = models.execute_kw(db, uid, password, 'llm.settings.line', 'read', [[llm_settings_line_id], ['display_name', 'model_name']])
llm_display_name = llm_line_data[0].get('display_name', 'Unknown') if llm_line_data else 'Unknown'
print(f"随机选择 llm_settings_line_id: {llm_settings_line_id} (name: {llm_display_name})")

# 搜索 Research Direction Prompt
rd_prompt_ids = models.execute_kw(db, uid, password, 'alpha.prompt.settings', 'search', [[['prompt_type', '=', 'Research Direction']]], {'limit': 1})
if not rd_prompt_ids:
    print("没有找到 prompt_type='Research Direction' 的 alpha.prompt.settings 记录")
    exit(1)
research_direction_prompt_id = rd_prompt_ids[0]
print(f"使用 Research Direction Prompt ID: {research_direction_prompt_id}")

# 搜索 Feature Engineering Prompt
fe_prompt_ids = models.execute_kw(db, uid, password, 'alpha.prompt.settings', 'search', [[['prompt_type', '=', 'Feature Engineering']]], {'limit': 1})
if not fe_prompt_ids:
    print("没有找到 prompt_type='Feature Engineering' 的 alpha.prompt.settings 记录")
    exit(1)
feature_engineering_prompt_id = fe_prompt_ids[0]
print(f"使用 Feature Engineering Prompt ID: {feature_engineering_prompt_id}")

# 创建研究方向记录
record_vals = {
    'region': region_id,
    'universe': universe_id,
    'llm_settings_line_id': llm_settings_line_id,
    'research_direction_prompt': research_direction_prompt_id,
    'feature_engineering_prompt': feature_engineering_prompt_id,
}

record_id = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'create', [record_vals])
print(f"成功创建研究方向记录，ID: {record_id}")

# 2. 调用 btn_quick_selection_datasets 按钮方法
result = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'btn_quick_selection_datasets', [[record_id]])
print(f"btn_quick_selection_datasets 执行结果: {result}")

# 验证数据集行是否已创建
datasets_line_ids = models.execute_kw(db, uid, password, 'alpha.research.datasets.line', 'search', [[['research_direction_id', '=', record_id]]])
print(f"已创建的数据集行数量: {len(datasets_line_ids)}")

# 获取数据集行的 dataset_code（数据集名字）
if datasets_line_ids:
    datasets_lines = models.execute_kw(db, uid, password, 'alpha.research.datasets.line', 'read', [datasets_line_ids, ['datasets_code']])
    dataset_codes = [line.get('datasets_code') for line in datasets_lines if line.get('datasets_code')]
    print(f"关联的数据集名字: {dataset_codes}")
    print(f"数据集总数: {len(dataset_codes)}")


# 3. 点击 btn_generate_research 按钮
print("\n开始生成研究方向...")
try:
    result = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'btn_generate_research', [[record_id]])
    print(f"btn_generate_research 执行结果: {result}")
except Fault as e:
    # 忽略 "cannot marshal None" 错误，因为方法返回 None 是正常的
    if "cannot marshal None" in str(e):
        print("btn_generate_research 已执行（返回 None，这是正常的）")
    else:
        raise

# 死循环等待 status = 'generated_research'
print("等待研究方向生成完成...")
while True:
    time.sleep(5)  # 每5秒检查一次
    record_data = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'read', [[record_id], ['status', 'message']])
    if not record_data:
        print("记录不存在")
        exit(1)
    current_status = record_data[0].get('status')
    message = record_data[0].get('message', '')
    print(f"当前状态: {current_status}, 消息: {message}")
    
    if current_status == 'generated_research':
        print("研究方向生成完成！")
        break
    elif current_status == 'failed':
        print(f"研究方向生成失败: {message}")
        exit(1)

# 4. 点击 btn_generate_prepare_feature_engineering_prompt 按钮
print("\n开始生成特征工程提示词...")
try:
    result = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'btn_generate_prepare_feature_engineering_prompt', [[record_id]])
    print(f"btn_generate_prepare_feature_engineering_prompt 执行结果: {result}")
except Fault as e:
    # 忽略 "cannot marshal None" 错误，因为方法返回 None 是正常的
    if "cannot marshal None" in str(e):
        print("btn_generate_prepare_feature_engineering_prompt 已执行（返回 None，这是正常的）")
    else:
        raise

# 5. 点击 btn_generate_feature 按钮
print("\n开始生成特征工程...")
try:
    result = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'btn_generate_feature', [[record_id]])
    print(f"btn_generate_feature 执行结果: {result}")
except Fault as e:
    # 忽略 "cannot marshal None" 错误，因为方法返回 None 是正常的
    if "cannot marshal None" in str(e):
        print("btn_generate_feature 已执行（返回 None，这是正常的）")
    else:
        raise

# 死循环等待 status = 'done'
print("等待特征工程生成完成...")
while True:
    time.sleep(5)  # 每5秒检查一次
    record_data = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'read', [[record_id], ['status', 'message']])
    if not record_data:
        print("记录不存在")
        exit(1)
    current_status = record_data[0].get('status')
    message = record_data[0].get('message', '')
    print(f"当前状态: {current_status}")
    
    if current_status == 'done':
        print("特征工程生成完成！")
        break
    elif current_status == 'failed':
        print(f"特征工程生成失败: {message}")
        exit(1)

# 6. 点击 btn_generate_idea 按钮
print("\n开始生成 Idea...")
try:
    result = models.execute_kw(db, uid, password, 'alpha.generate.research.direction', 'btn_generate_idea', [[record_id]])
    print(f"btn_generate_idea 执行结果: {result}")
except Fault as e:
    # 忽略 "cannot marshal None" 错误，因为方法返回 None 是正常的
    if "cannot marshal None" in str(e):
        print("btn_generate_idea 已执行（返回 None，这是正常的）")
    else:
        raise

print("\n全部流程执行完成！收工！")