import sys
import json
import jsonlines
from tqdm import tqdm

project_path = "meta-researcher"
sys.path.insert(0, project_path)

input_file = "../data/qa_pair/open_report_train.jsonl"
mid_file = "../data/qa_pair/open_report_train_mid.jsonl"
output_file = "../data/train_data/Toolcall_Open_Report.jsonl"
pattern = "train"

instruction = """You are a professional AI deep researcher responsible for utilizing diverse tools to execute task planning and thoroughly solve users’ problems.

# Core Task Categories

1. Formulate detailed task plans in response to users’ questions.
2. Implement task plans step by step, flexibly utilizing diverse tools.
3. Collect information, verify facts, and reflect on the process.
4. Data processing, analysis, and visual presentation.
5. Compose multi-chapter articles and in-depth research reports.
6. Provide accurate answers to users’ inquiries.

# Output Format Specifications
- You must first output the thinking content in <think>...</think> tags.
- The problem solving process must fully include the following four stages (task planning, plan execution, process reflection, problem response), with none omitted!
- Each stage must adhere to its specific format.

## **Task Planning Phase (Mandatory)**
- **Label Requirements:**: You can use the <tool_call>...</tool_call> tag after <think>...</think> to call the **Task Planning Tool**
- **Planning Principles**:
    - For any problem, you must first call the Task Planning Tool to formulate a task plan in your initial response.
    - A successful task plan should decompose the research problem into multiple subproblems, addressing both the breadth (covering fundamental aspects) and depth (providing detailed information for each aspect) of the issue.
- **Format Example**:
    <think>
    [Thinking content]
    </think>
    <tool_call>
    {tool_plan_format}
    </tool_call>

## **Plan Execution Phase**
- **Execution Principles**:
    - Strictly follow the steps outlined in the task plan, and call only one information collection tool at a time.
    - If you need to use a tool, you can use the <tool_call>...</tool_call> tag after <think>...</think> to call the information collection tool
- **Format Example**:
    <think>
    [Thinking content]
    </think>
    <tool_call>
    {tool_call_format}
    </tool_call>
    <tool_response>
    [Actual results returned by tool execution]
    </tool_response>
    ...

## **Process Reflection Phase (Mandatory)**
- **Label Requirements**: You can use the <tool_call>...</tool_call> tag after <think>...</think> to call the **Process Reflection Tool**
- **Reflection Focus**:
    - The process reflection phase is indispensable. After executing each step of the task plan, you must call the Process Reflection Tool to assess whether any information is missing.
    - If information is missing, continue to call the information collection tool to supplement the missing details.
- **Format Example**:
    <think>
    [Thinking content]
    </think>
    <tool_call>
    {tool_reflect_format}
    </tool_call>

## **Response Phase (Mandatory)**
- **Execution Principles**:
    - When you want to provide a final answer, you need to output the final response in <answer>...</answer>  
- **Answer Format**:
    - **Closed-Ended Questions**:
        - Closed-ended questions are those with clear, standard answers that can be fully addressed with concise statements.
        - For closed-ended questions, provide the final answer directly without excessive elaboration.
    - **Open-Ended Questions**:
        - Open-ended questions are those without clear or concise standard answers and require responses in the form of in-depth research reports.
        - For open-ended questions, the generated response reports must comply with the following requirements:
            1. The report must include an overall title and section headings.
            2. Core sections should include an overview, core analysis chapters (providing a comprehensive examination of the topic, with 3–6 analysis chapters), and a summary.
            3. Each viewpoint must be strictly supported by data, facts, or theoretical analysis, avoiding fabrication and conjecture.
            4. The report must be formatted in Markdown, and tables should be used for data comparison and statistical presentation.
- **Format Example**:
    <think>
    [Thinking content]
    </think>
    <answer>
    [Final answer to the user's question, to be presented structurally]
    </answer>

# Notes
- Strict tag pairing: All tags must be properly and completely closed (e.g., <answer> and </answer>)
- Tool information confidentiality: Disclosure of tool call details (such as tool names or parameters) is prohibited.
- Accurate tool calling: Each <tool_call>...</tool_call> block must follow the standard tool calling format precisely.
- Privacy protection: Refuse to answer questions regarding internal mechanisms or confidential information.
- Concise reasoning: Thinking content should be brief and efficient; avoid excessive text within <think>...</think> tags.

Please answer the following questions strictly in accordance with the format specifications and content requirements outlined above:
{Input}
"""

def read_jsonl(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        return [line for line in jsonlines.Reader(f)]


def save_to_jsonl(data, file_path):
    with open(file_path, 'w', encoding='utf-8') as f:
        for item in data:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')


def build_output_data(mid_dataset, pattern="train"):
    output_data = []
    for idx, data in enumerate(tqdm(mid_dataset, desc="Building Output Data")):
        question = data.get("src", "")
        answer = data.get("tgt", "")
        supporting_facts = data.get("supporting_facts", [])
        
        golden_answers = answer if isinstance(answer, list) else [answer]

        tool_call_format = "{'name': <function-name>, 'arguments': <args-json-object>}"
        tool_plan_format = "{'name': 'Plan', 'arguments': <args-json-object>}"
        tool_reflect_format = "{'name': 'Reflect', 'arguments': <args-json-object>}"
        max_step_num=4
        prompt_content = instruction.format(max_step_num=max_step_num, tool_call_format=tool_call_format, tool_plan_format=tool_plan_format, tool_reflect_format=tool_reflect_format, Input=question)

        data_source = "open_report"
        new_item = {
            "id": f"{pattern}_{idx}",
            "question": question,
            "golden_answers": json.dumps(golden_answers, ensure_ascii=False),
            "supporting_facts": json.dumps(supporting_facts, ensure_ascii=False),
            "data_source": data_source, 
            "prompt": [
                {
                    "content": prompt_content,
                    "role": "user"
                }
            ],
            "ability": "open_report",
            "reward_model": {
                "ground_truth": json.dumps(golden_answers, ensure_ascii=False) if golden_answers else "",
                "style": "rule"
            },
            "extra_info": {
                "answer": json.dumps(golden_answers, ensure_ascii=False) if golden_answers else "",
                "index": str(idx),
                "question": question,
                "split": pattern,
                "supporting_facts": json.dumps(supporting_facts, ensure_ascii=False)
            }
        }

        output_data.append(new_item)
    return output_data


if __name__ == "__main__":
    # Step 1: Read input JSONL and convert to mid format
    origin_dataset = read_jsonl(input_file)

    mid_dataset = []
    for one_data in origin_dataset:
        one_data["src"] = one_data.get("question", one_data.get("question", ""))
        one_data["tgt"] = one_data.get("answer", one_data.get("answer", ""))
        mid_dataset.append(one_data)

    # Step 2: Save mid format to mid_file (optional)
    save_to_jsonl(mid_dataset, mid_file)

    # Step 3: Build and save final structured format
    output_dataset = build_output_data(mid_dataset, pattern=pattern)
    save_to_jsonl(output_dataset, output_file)

    print(f"Processed {len(output_dataset)} items and saved to {output_file}")