# -*- coding: utf-8 -*-

import requests
# from load_dataset import load_doi
import json
import os
from datetime import datetime

API_KEY = 'app-YPSbl7Y2QMI755iQHbBZy7qP'
API_URL = 'http://localhost/v1/workflows/run'  # 按实际端口和路径填写

headers = {
    'Authorization': f'Bearer {API_KEY}',
    'Content-Type': 'application/json'
}
def load_doi():
    # 假设文件名为 arxiv-metadata.json
    file_path = "templates/dataset/arxiv.json"

    count = 0
    max_count = 100

    doi_list = []

    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            if count >= max_count:
                break
            data = json.loads(line)
            # 以你截图的字段为例，提取部分信息
            arxiv_id = data.get("id")
            title = data.get("title")
            authors = data.get("authors")
            abstract = data.get("abstract")
            categories = data.get("categories")
            doi_list.append(arxiv_id)
            count += 1
    return doi_list

from tqdm import tqdm

def polish_paper():

    doi_list = load_doi()
    
    # 创建输出目录
    output_dir = 'output/polish'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    for doi in tqdm(doi_list):
        url = f"https://arxiv.org/abs/{doi}"
        # print(url)
        # 如果你的应用有自定义 inputs，可以在这里填写
        inputs = {
            "url": url,
            # "input2": "value2",
        }

        payload = {
            "inputs": inputs,
            "response_mode": "blocking",  # 或 "streaming"
            "user": "abc-123",            # 可选，用户ID
            "conversation_id": ""        # 可选，首次对话留空，后续对话填前一次返回的 conversation_id
        }

        response = requests.post(API_URL, headers=headers, json=payload)

        if response.status_code == 200:
            print("运行成功，原始返回内容如下：")
            data = json.loads(response.text)
            print(data)  # 打印润色结果
            # 将润色结果保存到文件
            output_file = os.path.join(output_dir, f"{doi}.txt")
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(str(data['data']['outputs']))
            print(f"润色结果已保存到 {output_file}")

polish_paper()




