import gradio as gr
import sys
import json
from tinydb import Query
from pathlib import Path
from langchain.prompts import PromptTemplate
from sqlmodel import Session, select
from schema import Dataset, DatasetItem, Q_A
from db.dataset_store import get_all_dataset

sys.path.append(str(Path(__file__).resolve().parent.parent))
from schema import APIProvider, LLMRequest, LLMResponse, DatasetItem
from db import save_dataset
from tools import call_openai_api, process_markdown_file, generate_json_example
from global_var import get_docs, get_prompt_store, get_sql_engine, get_datasets, get_workdir

def dataset_generate_page():
    with gr.Blocks() as demo:
        gr.Markdown("## 数据集生成")
        with gr.Row():
            with gr.Column(scale=1):
                docs_list = [str(doc.name) for doc in get_docs()]
                initial_doc = docs_list[0] if docs_list else None
                prompts = get_prompt_store().all()
                prompt_list = [f"{p['id']} {p['name']}" for p in prompts]
                initial_prompt = prompt_list[0] if prompt_list else None
                
                # 初始化Dataframe的值
                initial_dataframe_value = []
                if initial_prompt:
                    selected_prompt_id = int(initial_prompt.split(" ")[0])
                    prompt_data = get_prompt_store().get(doc_id=selected_prompt_id)
                    prompt_content = prompt_data["content"]
                    prompt_template = PromptTemplate.from_template(prompt_content)
                    input_variables = prompt_template.input_variables
                    input_variables.remove("document_slice")
                    initial_dataframe_value = [[var, ""] for var in input_variables]
                # 从数据库获取API Provider列表
                with Session(get_sql_engine()) as session:
                    providers = session.exec(select(APIProvider)).all()
                    api_list = [f"{p.id} {p.model_id}" for p in providers]
                initial_api = api_list[0] if api_list else None
                
                api_dropdown = gr.Dropdown(
                    choices=api_list,
                    value=initial_api,
                    label="选择API",
                    interactive=True
                )
                doc_dropdown = gr.Dropdown(
                    choices=docs_list,
                    value=initial_doc,
                    label="选择文档",
                    interactive=True
                )
                prompt_dropdown = gr.Dropdown(
                    choices=prompt_list,
                    value=initial_prompt,
                    label="选择模板",
                    interactive=True
                )
                rounds_input = gr.Number(
                    value=1,
                    label="生成轮次",
                    minimum=1,
                    maximum=100,
                    step=1,
                    interactive=True
                )
                concurrency_input = gr.Number(
                    value=1,
                    label="并发数",
                    minimum=1,
                    maximum=10,
                    step=1,
                    interactive=True,
                    visible=False
                )
                dataset_name_input = gr.Textbox(
                    label="数据集名称",
                    placeholder="输入数据集保存名称",
                    interactive=True
                )
                prompt_choice = gr.State(value=initial_prompt)
                generate_button = gr.Button("生成数据集",variant="primary")
                doc_choice = gr.State(value=initial_doc)
                output_text = gr.Textbox(label="生成结果", interactive=False)
                api_choice = gr.State(value=initial_api)
            with gr.Column(scale=2):
                variables_dataframe = gr.Dataframe(
                    headers=["变量名", "变量值"],
                    datatype=["str", "str"],
                    interactive=True,
                    label="变量列表",
                    value=initial_dataframe_value  # 设置初始化数据
                )
        def on_doc_change(selected_doc):
            return selected_doc
        
        def on_api_change(selected_api):
            return selected_api
            
        def on_prompt_change(selected_prompt):
            if not selected_prompt:
                return None, []
            selected_prompt_id = int(selected_prompt.split(" ")[0])
            prompt_data = get_prompt_store().get(doc_id=selected_prompt_id)
            prompt_content = prompt_data["content"]
            prompt_template = PromptTemplate.from_template(prompt_content)
            input_variables = prompt_template.input_variables
            input_variables.remove("document_slice")
            dataframe_value = [] if input_variables is None else input_variables
            dataframe_value = [[var, ""] for var in input_variables]
            return selected_prompt, dataframe_value

        def on_generate_click(doc_state, prompt_state, api_state, variables_dataframe, rounds, concurrency, dataset_name, progress=gr.Progress()):
            dataset_db = get_datasets()
            if not dataset_db.search(Query().name == dataset_name):
                raise gr.Error("数据集名称已存在")
            
            doc = [i for i in get_docs() if i.name == doc_state][0]
            doc_files = doc.markdown_files
            document_slice_list = [process_markdown_file(doc) for doc in doc_files]
            prompt = [i for i in get_prompt_store().all() if i["id"] == int(prompt_state.split(" ")[0])][0]
            prompt = PromptTemplate.from_template(prompt["content"])
            with Session(get_sql_engine()) as session:
                api_provider = session.exec(select(APIProvider).where(APIProvider.id == int(api_state.split(" ")[0]))).first()

            variables_dict = {}
            for _, row in variables_dataframe.iterrows():
                var_name = row['变量名'].strip()
                var_value = row['变量值'].strip()
                if var_name:
                    variables_dict[var_name] = var_value

            prompt = prompt.partial(**variables_dict)
            
            dataset = Dataset(
                name=dataset_name,
                model_id=[api_provider.model_id],
                source_doc=doc,
                dataset_items=[]
            )
            
            total_slices = len(document_slice_list)
            for i, document_slice in enumerate(document_slice_list):
                progress((i + 1) / total_slices, desc=f"处理文档片段 {i + 1}/{total_slices}")
                request = LLMRequest(api_provider=api_provider,
                                    prompt=prompt.format(document_slice=document_slice),
                                    format=generate_json_example(DatasetItem))
                call_openai_api(request, rounds)

                for resp in request.response:
                    try:
                        content = json.loads(resp.content)
                        dataset_item = DatasetItem(
                            message=[Q_A(
                                question=content.get("question", ""),
                                answer=content.get("answer", "")
                            )]
                        )
                        dataset.dataset_items.append(dataset_item)
                    except json.JSONDecodeError as e:
                        print(f"Failed to parse response: {e}")

            # 保存数据集到TinyDB
            dataset_db.insert(dataset.model_dump())

            save_dataset(dataset_db,get_workdir(),dataset_name)
            
            return f"数据集 {dataset_name} 生成完成，共 {len(dataset.dataset_items)} 条数据"

        doc_dropdown.change(on_doc_change, inputs=doc_dropdown, outputs=doc_choice)
        prompt_dropdown.change(on_prompt_change, inputs=prompt_dropdown, outputs=[prompt_choice, variables_dataframe])
        api_dropdown.change(on_api_change, inputs=api_dropdown, outputs=api_choice)

        generate_button.click(
            on_generate_click,
            inputs=[doc_choice, prompt_choice, api_choice, variables_dataframe, rounds_input, concurrency_input, dataset_name_input],
            outputs=output_text
        )

    return demo

if __name__ == "__main__":
    from global_var import init_global_var
    init_global_var("workdir")
    demo = dataset_generate_page()
    demo.launch()