import json
import os
import re
import time
from typing import Any, Optional, Union

import jinja2
from pydantic import BaseModel, Field, PrivateAttr
from sglang import function, system, user, assistant, gen, set_default_backend, OpenAI
from sglang.lang.interpreter import ProgramState, SglGen


prompt_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "prompts")


def check_is_leaf_project(directory: str):
    if not os.path.exists(os.path.join(directory, ".leaf")):
        raise ValueError(f"{directory} must be a project directory that generated by leaf-playground CLI")


def add_and_show_message(
    state: ProgramState, msg_wrapper: Union[system, user, assistant], msg: Union[str, SglGen], show: bool = True
):
    if isinstance(msg, SglGen) and not msg.name:
        raise ValueError("Must set name to SglGen instance")
    state += msg_wrapper(msg)
    if msg_wrapper != system and show:
        if isinstance(msg, str):
            message = state.messages()[-1]
            print(f"{message['role'].title()}: {message['content']}", end="\n\n")
        else:
            print(f"{'User' if msg_wrapper == user else 'Assistant'}: ", end="")
            for text in state.text_iter(var_name=msg.name):
                print(text, end="", flush=True)
            print()


def load_prompt(*rel_path: str):
    with open(os.path.join(prompt_dir, *rel_path), "r", encoding="utf-8") as f:
        return f.read()


def load_code(reference_file: Optional[str], default: Optional[str]):
    if not reference_file:
        return default
    with open(reference_file, "r", encoding="utf-8") as f:
        return f.read()


def extract_code_and_save(code: str, file_name: str, save_dir: str):
    pattern = r"```python\n# " + file_name + r".py(.*?)```"
    match = re.search(pattern, code, re.DOTALL)
    if match:
        code = match.group(1).strip()

    with open(os.path.join(save_dir, f"{file_name}.py"), "w", encoding="utf-8") as f:
        f.write(code)


class PipelineConfig(BaseModel):
    target_project: str = Field(default=...)
    reference_project: Optional[str] = Field(default=None)
    openai_api_key: Optional[str] = Field(default=None)

    definition_completion: bool = Field(default=True)
    agent_completion: bool = Field(default=True)
    scene_completion: bool = Field(default=True)

    _target_project_config: dict = PrivateAttr(default=None)
    _target_pkg: str = PrivateAttr(default=None)
    _reference_pkg: str = PrivateAttr(default=None)

    def model_post_init(self, __context: Any) -> None:
        if self.openai_api_key is None:
            self.openai_api_key = os.environ.get("OPENAI_API_KEY", None)
        if not self.openai_api_key:
            raise ValueError(
                "Must provide the `openai_api_key` argument or the `OPENAI_API_KEY` environment variable。"
            )

        check_is_leaf_project(self.target_project)
        if self.reference_project:
            check_is_leaf_project(self.reference_project)

        self._target_project_config = self.load_project_config(self.target_project)
        self._target_pkg = self.get_pkg_path(self.target_project)
        if self.reference_project:
            self._reference_pkg = self.get_pkg_path(self.reference_project)

    @staticmethod
    def load_project_config(project_path: str):
        with open(os.path.join(project_path, ".leaf", "project_config.json"), "r", encoding="utf-8") as f:
            proj_config = json.load(f)
        return proj_config

    @staticmethod
    def get_pkg_path(project_path: str):
        proj_config = PipelineConfig.load_project_config(project_path)
        return os.path.join(project_path, proj_config["name"])

    @property
    def target_project_config(self):
        return self._target_project_config

    @property
    def target_pkg_path(self):
        return self._target_pkg

    @property
    def reference_pkg_path(self):
        return self._reference_pkg


class Pipeline:
    def __init__(self, config: PipelineConfig):
        set_default_backend(OpenAI("gpt-4-1106-preview", api_key=config.openai_api_key))

        self.config = config

    def _definition_completion(self):
        initial_msg = jinja2.Template(load_prompt("system", "product_manager.prompt")).render(
            leaf_playground_overview=load_prompt("docs", "overview.prompt"),
            definition_doc=jinja2.Template(load_prompt("docs", "definition.prompt")).render(
                data_message_module_reference=load_prompt("docs", "references", "data_message_module.prompt"),
                data_environment_module_reference=load_prompt("docs", "references", "data_environment_module.prompt"),
                core_definition_module_reference=load_prompt("docs", "references", "core_definition_module.prompt"),
            ),
            scene_definition_example=load_code(
                reference_file=(
                    None
                    if not self.config.reference_pkg_path
                    else os.path.join(self.config.reference_pkg_path, "scene_definition.py")
                ),
                default=None if self.config.reference_pkg_path else load_prompt("examples", "scene_definition.py"),
            ),
            project_name=self.config.target_project_config["name"],
            project_description=input(
                jinja2.Template(load_prompt("user_inputs", "ask_for_project_description.prompt")).render(
                    project_name=self.config.target_project_config["name"]
                )
            ),
        )

        _finished = False

        @function
        def _complete(s, initial_msg):
            nonlocal _finished

            add_and_show_message(s, user, initial_msg, False)
            system_msg = "阶段：第一阶段；要求：{0}"
            while True:
                add_and_show_message(
                    s,
                    system,
                    system_msg.format(
                        "先响应用户的回复，然后一步步地思考（think step by step）并使用自然语言将需求细化，暂时不要提问"
                    ),
                )
                add_and_show_message(s, assistant, gen("detailed_solution", max_tokens=4096))
                add_and_show_message(
                    s,
                    system,
                    system_msg.format("判断是否需要向用户提问或寻求改进意见，直接回答“是”或“否”，不要说其他的话"),
                )
                add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                if "否" in s.get_var("judgement"):
                    break
                add_and_show_message(
                    s,
                    system,
                    system_msg.format(
                        "先告知用户你接下来要做的事情，然后分点列出你需要向用户提问的全部问题或向用户寻求改进意见"
                    ),
                )
                add_and_show_message(s, assistant, gen("questions", max_tokens=128))
                add_and_show_message(s, user, input())
                add_and_show_message(
                    s,
                    system,
                    system_msg.format("根据用户回复判断是否需要修改需求细节，直接回答“是”或“否”，不要说其他的话"),
                )
                add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                if "否" in s.get_var("judgement"):
                    break
            system_msg = "阶段：第二阶段；要求：{0}"
            while True:
                add_and_show_message(
                    s,
                    system,
                    system_msg.format(
                        "结合用户在一开始给你的指南并根据最新细化的需求，生成 scene_definition.py，然后询问用户是否有需要补充或修改的"
                    ),
                )
                add_and_show_message(s, assistant, gen("scene_definition.py", max_tokens=4096))
                add_and_show_message(s, user, input())
                add_and_show_message(
                    s,
                    system,
                    "根据用户的回复判断你是否需要对 scene_definition.py 进行修改？直接回答“是”或“否”，不要说其他的话",
                )
                add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                if "否" in s.get_var("judgement"):
                    break

            extract_code_and_save(s.get_var("scene_definition.py"), "scene_definition", self.config.target_pkg_path)

            _finished = True

        _complete.run(initial_msg, temperature=0.7, stream=True)

        while not _finished:
            time.sleep(0.01)

    def _agent_completion(self):
        target_scene_definition = load_code(os.path.join(self.config.target_pkg_path, "scene_definition.py"), None)
        if not target_scene_definition:
            raise NotImplementedError("must implement scene_definition.py before complete agents' code")

        initial_msg = jinja2.Template(load_prompt("system", "base_agent_developer.prompt")).render(
            leaf_playground_overview=load_prompt("docs", "overview.prompt"),
            domain_concept=load_prompt("docs", "domain_concept.prompt"),
            base_agent_dev_doc=jinja2.Template(load_prompt("docs", "base_agent_dev.prompt")).render(
                core_scene_agent_module_reference=load_prompt("docs", "references", "core_scene_agent_module.prompt"),
                data_message_module_reference=load_prompt("docs", "references", "data_message_module.prompt"),
            ),
            static_agent_module_name="examiner",
            static_agent_example_code=load_prompt("examples", "base_agents", "examiner.py"),
            dynamic_agent_module_name="examinee",
            dynamic_agent_example_code=load_prompt("examples", "base_agents", "examinee.py"),
            scene_definition_code=target_scene_definition,
        )

        _finished = False

        @function
        def _complete(s, initial_msg):
            nonlocal _finished

            add_and_show_message(s, user, initial_msg, False)
            add_and_show_message(
                s, system, "请你首先根据 scene_definition.py 的信息提取出所有角色名，并判断他们是动态还是静态的角色"
            )
            add_and_show_message(s, assistant, gen("retrieve_role_info", max_tokens=512, temperature=0.0), False)
            add_and_show_message(s, system, "接下来，你需要依次生成各角色的智能体模块代码")
            while True:
                add_and_show_message(
                    s, system, "请你明确现在你要生成的是哪个角色的智能体模块代码？直接给我角色名，不要说其他的话。"
                )
                add_and_show_message(s, assistant, gen("role_name", max_tokens=16, temperature=0.1), False)
                role_name = s.get_var("role_name")
                add_and_show_message(
                    s,
                    system,
                    "首先告知用户你接下来要生成的角色是哪个，然后另起一行，根据用户指定的格式，结合"
                    f" scene_definition.py 中对于角色 {role_name} 的定义，"
                    "依据用户给你的指南并参考示例代码，生成该角色的智能体模块代码，直接给出代码，并在代码生成完毕后向用户征询改进建议",
                )
                add_and_show_message(s, assistant, gen(f"{role_name}.py", max_tokens=4096, temperature=0.5))
                while True:
                    add_and_show_message(s, user, input())
                    add_and_show_message(
                        s, system, "请你根据用户的反馈判断是否需要修改代码实现？直接回答“是”或“否”，不要说其他的话"
                    )
                    add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                    if "否" in s.get_var("judgement"):
                        break
                    add_and_show_message(
                        s,
                        system,
                        f"请你结合用户的反馈修改 {role_name}.py 代码，直接给出代码，并在代码生成完毕后再次向用户征询改进建议",
                    )
                    add_and_show_message(s, assistant, gen(f"{role_name}.py", max_tokens=4096, temperature=0.5))
                extract_code_and_save(
                    s.get_var(f"{role_name}.py"), role_name, os.path.join(self.config.target_pkg_path, "agents")
                )
                add_and_show_message(
                    s,
                    system,
                    "请你判断你是否已经生成完 scene_definition.py"
                    " 中定义的全部角色的智能体模块代码？直接回答“是”或“否”，不要说其他的话",
                )
                add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                if "是" in s.get_var("judgement"):
                    break

            _finished = True

        _complete.run(initial_msg, stream=True)

        while not _finished:
            time.sleep(0.01)

    def _scene_completion(self):
        target_scene_definition = load_code(os.path.join(self.config.target_pkg_path, "scene_definition.py"), None)
        if not target_scene_definition:
            raise NotImplementedError("must implement scene_definition.py before complete agents' code")

        agent_modules = []
        for root, dirs, files in os.walk(os.path.join(self.config.target_pkg_path, "agents")):
            for file in files:
                if file.endswith(".py") and "__init__" not in file:
                    agent_modules.append(os.path.join(root, file))
        if not agent_modules:
            raise NotImplementedError("must implement all agents base class that defined in scene_definition.py")

        initial_msg = jinja2.Template(load_prompt("system", "scene_developer.prompt")).render(
            leaf_playground_overview=load_prompt("docs", "overview.prompt"),
            domain_concept=load_prompt("docs", "domain_concept.prompt"),
            scene_dev_doc=jinja2.Template(load_prompt("docs", "scene.prompt")).render(
                core_scene_module_reference=load_prompt("docs", "references", "core_scene_module.prompt"),
                data_log_module_reference=load_prompt("docs", "references", "data_log_module.prompt"),
            ),
            scene_example_code=load_code(
                reference_file=(
                    None
                    if not self.config.reference_pkg_path
                    else os.path.join(self.config.reference_pkg_path, "scene.py")
                ),
                default=None if self.config.reference_pkg_path else load_prompt("examples", "scene.py"),
            ),
            scene_definition_code=target_scene_definition,
            agents_code="\n\n".join([
                f"```python\n# {agent_module}\n\n{load_code(agent_module, None)}\n```"
                for agent_module in agent_modules
            ]),
        )

        _finished = False

        @function
        def _complete(s, initial_msg):
            nonlocal _finished

            add_and_show_message(s, user, initial_msg, False)
            add_and_show_message(
                s,
                system,
                "请你首先根据 scene_definition.py 的信息，自行描述一遍完整的场景运行流程，并询问用户是否需要调整",
            )
            add_and_show_message(s, assistant, gen("scene_flow_desc", max_tokens=4096, temperature=0.5))
            while True:
                add_and_show_message(s, user, input())
                add_and_show_message(
                    s,
                    system,
                    "请你根据用户的反馈判断是否需要调整场景运行流程的描述？直接回答“是”或“否”，不要说其他的话",
                )
                add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                if "否" in s.get_var("judgement"):
                    break
                add_and_show_message(
                    s, system, "请你结合用户的反馈调整场景运行流程的描述，并再次询问用户是否需要进一步地调整"
                )
                add_and_show_message(s, assistant, gen("scene_flow_desc", max_tokens=4096, temperature=0.5))
            add_and_show_message(
                s,
                system,
                f"请你根据用户指定的格式，结合上面的信息，依据用户给你的指南并参考示例代码，生成 scene.py 代码，"
                f"直接给出代码，并在代码生成完毕后向用户征询改进建议",
            )
            add_and_show_message(s, assistant, gen("scene.py", max_tokens=4096, temperature=0.5))
            while True:
                add_and_show_message(s, user, input())
                add_and_show_message(
                    s, system, "请你根据用户的反馈判断是否需要修改代码实现？直接回答“是”或“否”，不要说其他的话"
                )
                add_and_show_message(s, assistant, gen("judgement", max_tokens=2, temperature=0.0), False)
                if "否" in s.get_var("judgement"):
                    break
                add_and_show_message(
                    s,
                    system,
                    "请你结合用户的反馈修改 scene.py 代码，直接给出代码，并在代码生成完毕后再次向用户征询改进建议",
                )
                add_and_show_message(s, assistant, gen("scene.py", max_tokens=4096, temperature=0.5))
            extract_code_and_save(s.get_var("scene.py"), "scene", self.config.target_pkg_path)

            _finished = True

        _complete.run(initial_msg, stream=True)

        while not _finished:
            time.sleep(0.01)

    def run(self):
        if self.config.definition_completion:
            self._definition_completion()
        if self.config.agent_completion:
            self._agent_completion()
        if self.config.scene_completion:
            self._scene_completion()


__all__ = ["Pipeline", "PipelineConfig"]
