import asyncio
from pathlib import Path

import pandas as pd
from dotenv import load_dotenv

# from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.settings import ModelSettings

from llm.config import load_config
from llm.factory import LLMModelProvider
from llm.schemas import DatasetRecord, Message, RobotPlan, SFTFormat

load_dotenv()


def save_to_csv(record: DatasetRecord,
                path: str | Path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    exists = path.exists()

    df = pd.DataFrame([record.model_dump()])
    df.to_csv(path,
              mode="a" if exists else "w",
              header=not exists,
              index=False,
              encoding="utf-8")

def save_to_jsonl(sft_format: SFTFormat,
                  path: str | Path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    
    with path.open("a", encoding="utf-8") as file:
        file.write(sft_format.model_dump_json())
        file.write("\n")


def record_to_sft(record: DatasetRecord) -> SFTFormat:
    return SFTFormat(
        message=[
            Message(
                role="system",
                content=record.system_prompt
            ),
            Message(
                role="user",
                content=record.user_prompt
            ),
            Message(
                role="assistant",
                content=record.assistant_content
            )
        ]
    )

def format_robot_plan(plan: RobotPlan):
    return "\n".join(
        f"{index}. {step}"
        for index, step in enumerate(plan.steps, start=1)
    )


async def main():
    config = load_config("./build-dataset/config.yml")
    model = await LLMModelProvider.create_from_config(config.llm)


    # model = await LLMModelProvider.create_codex(model_name='gpt-5.6-luna')
    # model = LLMModelProvider.create_openai_compatible(model_name="qwen/qwen3-8b",
    #                                                   api_key=os.getenv("API_KEY"),
    #                                                   base_url=os.getenv("BASE_URL"))
    agent = Agent(model=model,
                  output_type=RobotPlan,
                  system_prompt=config.prompts.system,
                  instructions=config.prompts.instructions)
    settings = ModelSettings(
        **config.generation.model_dump()
    )

    task="Перемести деталь A из точки PICK_A в точку PLACE_A. Безопасная высота подъема — 200 мм."
    user_prompt = config.prompts.user.format(task=task)
    result = await agent.run(user_prompt,
                             model_settings=settings)

    assistant_content = format_robot_plan(result.output)
    record = DatasetRecord(
        task=task,
        system_prompt=config.prompts.system,
        instructions=config.prompts.instructions,
        user_prompt=user_prompt,
        model_response=result.output.model_dump_json(),
        assistant_content=assistant_content,
        provider=config.llm.provider,
        model=config.llm.model,
        temperature=config.generation.temperature,
        max_tokens=config.generation.max_tokens
    )

    save_to_csv(record=record,
                path="./dataset/robot-move.csv")
    sft_format = record_to_sft(record)
    save_to_jsonl(sft_format=sft_format,
                  path="./dataset/robot-move.jsonl")


if __name__ == "__main__":
    asyncio.run(main())