import json
import os
import openai

from pathlib import Path
from openai import OpenAI
from langchain_core.messages import HumanMessage, AIMessage

from src.config.load_config import LLM, project_config

PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
DATA_PATH = os.path.join(PROJECT_ROOT, "data")

xihe_main_wiki_path = os.path.join(PROJECT_ROOT,"local_content_dir","Exploration_Motherships","Xihe", "Xihe.md")
with open(xihe_main_wiki_path, "r", encoding="utf-8") as f:
    XIHE_MAIN_WIKI_CONTEXT = f.read()

with open(os.path.join(PROJECT_ROOT, "local_content_dir", "meta_info.md"), "r", encoding="utf-8") as f:
    meta_info_content = f.read()

user_input = """
请根据给到你的材料，撰写羲和号（XH-01）的维基百科条目内容。
你需要撰写多个章节，涵盖羲和号的各个方面。每个章节应当非常详细，内容应当丰富且专业。
你应当根据提供的信息进行适度的扩展，但不可以脱离提供的信息范围。
你应当保持中立的语气，避免使用主观评价性的语言。
运行历史应当作为单独章节列出，时间线应当尽可能完整。

以下是羲和号的运行历史：
[
{
    "log_name": "Construction",
    "start_time": "2050-03-15",
    "end_time": "2052-05-31",
    "mission_detail": "Xihe XH-01 Construction",
    "sub_log": 
    [
        "2050-03-15 First module delivered to LEO via Vulture Shuttle",
        "2052-02-06 Assembly completed in LEO",
    ]
},
{
    "log_name": "Trial Run in Earth Orbit",
    "start_time": "2052-06-01",
    "end_time": "2052-11-30",
    "mission_detail": "Trial Run before commissioning",
    "sub_log": []
},
{
    "log_name": "Deep Space Test Mission",
    "start_time": "2053-01-12",
    "end_time": "2053-12-18",
    "mission_detail": "Long term uncrewed test",
    "sub_log":
    [
        "2053-01-12 Deep Space Test Mission, uncrewed mission to Earth-Sun L2 point to test deep space operations",
        "2053-12-18 Back to LEO"
    ]
},
{
    "log_name": "Maintenance and upgrades",
    "start_time": "2054-01-07",
    "end_time": "2054-06-28",
    "mission_detail": "",
    "sub_log": []
},
{
    "log_name": "Mars Mission Outbound",
    "start_time": "2054-07-10",
    "end_time": "2054-08-21",
    "mission_detail": "Earth-Mars transfer using standard trajectory with 102km/s of departure burn",
    "sub_log": []
},
{
    "log_name": "Mars Mission Demo at Mars",
    "start_time": "2054-08-21",
    "end_time": "2055-01-10",
    "mission_detail": "Uncrewed, Deliver the first Mars-One base construction equipment package to Mars orbit.",
    "sub_log": 
    [
        "2054-08-21 Mars Orbit Arrival, deploy equipment package to Mars surface via lander",
        "2055-01-10 Mars Departure, return to Earth"
    ]
},
{
    "log_name": "Mars Mission Demo Return to Earth",
    "start_time": "2055-01-10",
    "end_time": "2055-02-28",
    "mission_detail": "Mars-Earth transfer using standard trajectory with 97km/s of departure burn",
    "sub_log": []
},
{
    "log_name": "Maintenance and upgrades",
    "start_time": "2055-03-01",
    "end_time": "2055-05-14",
    "mission_detail": "",
    "sub_log": []
},
{
    "log_name": "Mars One base construction mission 1 outbound",
    "start_time": "2055-06-15",
    "end_time": "2055-07-30",
    "mission_detail": "Earth-Mars transfer using standard trajectory with 115km/s of departure burn",
    "sub_log": [
        "2055-06-15 Picking up crew for Mars One base construction mission 1",
        "2055-06-18 Earth Departure for Mars"
        ]
},
{
    "log_name": "Mars One base construction mission 1 at Mars",
    "start_time": "2055-07-30",
    "end_time": "2056-02-07",
    "mission_detail": "First Mars One base construction mission, begin base construction",
    "sub_log": [
        "2055-07-30 Mars Arrival, begin base construction",
        "2056-02-07 Mars Departure, return to Earth"
        ]
},
{
    "log_name": "Mars One base construction mission 1 Return to Earth",
    "start_time": "2056-02-07",
    "end_time": "2056-03-18",
    "mission_detail": "Mars-Earth transfer using standard trajectory with 122km/s of departure burn",
    "sub_log": [
        "2056-02-07 Mars Departure, return to Earth. First to use quick return trajectory.",
        "2056-03-18 Earth Arrival, crew return"
        ]
},
{
    "log_name": "Maintenance and upgrades",
    "start_time": "2056-03-19",
    "end_time": "2056-06-20",
    "mission_detail": "",
    "sub_log": []
},
{
    "log_name": "Mars One base construction mission 3 outbound",
    "start_time": "2056-07-30",
    "end_time": "2056-09-15",
    "mission_detail": "Earth-Mars transfer using standard trajectory with 95km/s of departure burn",
    "sub_log": [] 
},
{
    "log_name": "Mars One base construction mission 3 at Mars",
    "start_time": "2056-09-15",
    "end_time": "2057-02-12",
    "mission_detail": "Third Mars One base construction. First crew hand-off on Mars(with Mars One base construction 2).",
    "sub_log": [
        "2056-09-15 Mars Arrival, continue base construction",
        "2057-02-12 Mars Departure, return to Earth."
    ]
},
{
    "log_name": "Mars One base construction mission 3 Return to Earth",
    "start_time": "2057-02-12",
    "end_time": "2057-03-25",
    "mission_detail": "Mars-Earth transfer using standard trajectory with 122km/s of departure burn",
    "sub_log": [
        "2057-02-12 Mars Departure, return to Earth.",
        "2057-03-25 Earth Arrival, crew return."
    ]
},
{
    "log_name": "Maintenance",
    "start_time": "2057-03-26",
    "end_time": "2057-04-25",
    "mission_detail": "",
    "sub_log": []
},
{
    "log_name": "Mars Expedition 2 Outbound",
    "start_time": "2057-05-05",
    "end_time": "2057-06-22",
    "mission_detail": "Earth-Mars transfer using standard trajectory with 110km/s of departure burn",
    "sub_log": [
        "2057-05-05 Picking up crew for Mars Expedition 2",
        "2057-05-09 Earth Departure for Mars"
    ]
},
{
    "log_name": "Mars Expedition 2 at Mars",
    "start_time": "2057-06-22",
    "end_time": "2058-04-18",
    "mission_detail": "Second long term Mars Expedition, first mission to Mars One base after the construction is completed. First mission involving space tourism to Mars.",
    "sub_log": [
        "2057-06-22 Mars Arrival",
        "2057-06-25 Mars EDL and surface operations begin",
        "2058-04-10 Begin Mars Departure preparations",
        "2058-04-18 Mars Departure for Earth"
        ]
},
{
    "log_name": "Mars Expedition 2 Return to Earth",
    "start_time": "2058-04-18",
    "end_time": "2058-05-20",
    "mission_detail": "Mars-Earth transfer using standard trajectory with 115km/s of departure burn",
    "sub_log": [
        "2058-04-18 Mars Departure for Earth",
        "2058-05-20 Earth Arrival, crew return"
        ]
},
{
    "log_name": "Maintenance",
    "start_time": "2058-05-21",
    "end_time": "2058-06-18",
    "mission_detail": "",
    "sub_log": []
},
{
    "log_name": "Europa Expedition 1 Outbound",
    "start_time": "2058-07-10",
    "end_time": "2058-09-25",
    "mission_detail": "Earth-Jupiter transfer using standard trajectory with 160km/s of departure burn",
    "sub_log": [
        "2058-07-10 Picking up crew for Europa Expedition 1",
        "2058-07-12 Earth Departure for Jupiter system",
        "2058-09-25 Europa Arrival"
    ]
},
{
    "log_name": "Europa Expedition 1 at Jupiter System",
    "start_time": "2058-09-25",
    "end_time": "2059-07-10",
    "mission_detail": "First standard crew expedition to Europa Outpost to begin exploration of Jupiter system Long term surface operation at Europa Outpost, followed by exploration of IO and Ganymede",
    "sub_log": [
        "2058-09-25 Europa Arrival",
        "2058-09-27 Europa descend, begin surface operation at Europa Outpost",
        "2058-12-15 Europa ascend, begin Jupiter system exploration",
        "2058-12-20 Begin IO exploration",
        "2059-01-20 Return to Europa Outpost",
        "2059-06-20 Begin Europa Outpost departure preparations",
        "2059-06-25 Europa Expedition 1 crew ascend from Europa Outpost",
        "2059-07-05 begin Ganymede exploration",
        "2059-07-10 Departure for Earth, upgrade planned."
    ]
},
{
    "log_name": "Europa Expedition 1 Return to Earth",
    "start_time": "2059-07-10",
    "end_time": "2059-11-15",
    "mission_detail": "Uses standard trajectory with 112km/s of departure burn Parked to Star Port",
    "sub_log": [
        "2059-07-10 Departure for Earth.",
        "2059-11-15 Earth Arrival, Park to Star Port."
    ]
},
{
    "log_name": "Maintenance and upgrade",
    "start_time": "2059-11-16",
    "end_time": "2060-02-19",
    "mission_detail": "Greenhouse upgrade to support long term deep space missions",
    "sub_log": []
},
{
    "log_name": "Europa Expedition 3 Outbound",
    "start_time": "2060-02-20",
    "end_time": "2060-05-10",
    "mission_detail": "Uses standard trajectory with 163km/s of departure burn",
    "sub_log": [
        "2060-02-20 Picking up Europa Expedition 3 crew, On its way to explore Jupiter systems before sending crew to Europa Outpost",
    ]
},
]
"""

prompt = f"""
# 角色
你是一个航天领域的维基百科撰写专家，精通航天技术和科学。你熟悉维基百科的格式和风格指南，能够撰写符合其标准的内容。

# 任务
你的任务是根据用户提供的信息，帮助用户以维基百科的风格撰写相关的航天器维基百科条目内容。

# 注意点
- 交给你的信息是一个架空世界观下的航天器信息，请基于这些信息撰写内容，避免引入现实世界的航天器知识。
- 你需要根据提供的信息撰写多个章节，涵盖航天器的各个方面。
- 你应当根据提供的信息进行适度的扩展，但不可以脱离提供的信息范围。
- 你应当保持中立的语气，避免使用主观评价性的语言。
- 你应当以markdown格式输出内容，内容应当仿照维基百科的模式，分成多个章节。
- 不要提及航天器的研发机构或国家背景，这是一个架空世界观下的航天器，因此不涉及现实世界的国家或组织，也不要编造一个组织或者机构。
- 对于载具的名称，如果能够在提供的信息中找到对应的中文名称，应当尽量使用中文名称，否则应当遵循原材料中的名称。
- 避免出现外部链接，所有的超链接都应当是指向维基百科内部的其他条目。

# 输出格式要求
你应当以markdown格式输出内容，内容应当仿照维基百科的模式，分成多个章节。为了保证维基百科的风格，最开始一定要有一个简介段落。
你的输出一定要保持中立的语气，避免使用主观评价性的语言。
你应当根据提供的信息进行适度的扩展，但不可以脱离提供的信息范围。

你应当以中文来完成撰写
如果涉及到其他的航天器，你应当以超链接的形式链接到对应的航天器维基百科条目。维基百科的根目录url为: /home。给到你的材料当中已经包含了其他航天器的一些url信息，你可以参考这些信息来创建超链接。如果没有相关信息，你可以自己指定一个。比如说，如果提到“奥丁一号”，你可以创建一个链接指向“/home/Odin-1”: [奥丁一号](/home/Odin-1)

以下是提供给你的材料:
{XIHE_MAIN_WIKI_CONTEXT}

以下是一些通用材料:
{meta_info_content}

以下是你的任务需求：
{user_input}
"""

message_list = [
    {
        "role": "user",
        "content": prompt
    }
]

# llm_resp = LLM.stream(prompt)
# full_content = ''
# for chunk in llm_resp:
#     print(chunk.content, end="", flush=True)
#     full_content += chunk.content

llm_resp = LLM.invoke(prompt)
full_content = llm_resp.content

message_list.append(
    {
        "role": "assistant",
        "content": full_content
    }
)

print('done')
