import asyncio

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, DatasetWriter
from llm.schemas import (RobotPlan, RobotDatasetRecord, 
                         Prompts, ModelResponse, 
                         ModelInfo, GenerationSettings)

load_dotenv()


def format_robot_plan(plan: RobotPlan):
    return "\n".join(
        f"{step}"
        for step in plan.steps
    )


async def generate_record(task: str,
                           agent: Agent,
                           settings: ModelSettings,
                           config,
                           semaphore: asyncio.Semaphore) -> RobotDatasetRecord:

    async with semaphore:
        user_prompt = config.prompts.user.format(task=task)
    
    result = await agent.run(user_prompt,
                             model_settings=settings)
    assistant_content = format_robot_plan(result.output)

    return RobotDatasetRecord(
        task=task,
        prompts=Prompts(
            system_prompt=config.prompts.system,
            instructions=config.prompts.instructions,
            user_prompt=user_prompt
        ),
        response=ModelResponse(
            model_response=result.output.model_dump_json(),
            assistant_content=assistant_content
        ),
        model_info=ModelInfo(
            provider=config.llm.provider,
            model=config.llm.model
        ),
        generation=GenerationSettings(
            temperature=config.generation.temperature,
            max_tokens=config.generation.max_tokens
        )
    )
    

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()
    )

    tasks = [
        (
            "Перемести деталь A из точки PICK_A в точку PLACE_A. "
            "Безопасная высота подъема — 200 мм."
        ),
        (
            "Перемести деталь B из точки PICK_B в точку PLACE_B. "
            "Безопасная высота подъема — 150 мм."
        ),
        (
            "Возьми крышку из точки COVER_PICK и перемести "
            "ее в точку COVER_PLACE."
        ),
    ]

    semaphore = asyncio.Semaphore(config.runtime.concurrency)

    records = await asyncio.gather(*[
        generate_record(
            task,
            agent,
            settings,
            config,
            semaphore,
        )
        for task in tasks
    ])

    for record in records:
        dataset = DatasetWriter(
            dataset_record=record,
            base_path="./dataset"
        )

        dataset.to_csv(filename="robot-plan.csv")
        dataset.to_jsonl(filename="robotplan.jsonl")


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