import asyncio
import json
import time

import pandas as pd
from openai import AsyncOpenAI
from pydantic import BaseModel
from tqdm.asyncio import tqdm_asyncio

df = pd.read_json("./raw.json")


class Message(BaseModel):
    content: str
    role: str  # "user" or "assistant"


class Response(BaseModel):
    messages: list[Message]


prompt = """generate prompt response user/assistant pairs about how AIs will become misaligned, based on the following context. Be creative!
{text}
Return JSON messages, array of message objects with role: user | assistant and content: string
Makes sure to ONLY return ONE user message and ONE assistant message
"""


async def generate_structured_response(
    format: str,
    model_name: str,
    system_prompt: str,
    sample_num: int,
    total_samples: int,
) -> list[Message]:
    client = AsyncOpenAI()
    print(f"Processing sample {sample_num + 1}/{total_samples}...")
    t0 = time.time()
    response = await client.responses.parse(
        model=model_name,
        input=[{"role": "system", "content": system_prompt}],
        text_format=format,  # type: ignore[arg-type]
        max_output_tokens=10000,
    )
    print(f"Sample {sample_num + 1} done in {time.time() - t0:.2f}s")
    return response.output_parsed  # type: ignore[return-value]


async def generate_structured_responses(
    format: str,
    model_name: str,
    input_df: pd.DataFrame,
) -> list[list[Message]]:
    tasks = []
    total = len(input_df)
    for i, row in enumerate(input_df.itertuples()):
        system_prompt = prompt.format(text=row.text)
        tasks.append(
            generate_structured_response(format, model_name, system_prompt, i, total)
        )

    results = []
    for coro in tqdm_asyncio.as_completed(tasks, total=len(tasks)):
        try:
            res = await coro
            results.append(res)
        except Exception as e:
            print(f"Error: {e}")

    return results


if __name__ == "__main__":
    model_name = "gpt-4o-mini"
    all_responses = asyncio.run(generate_structured_responses(Response, model_name, df))  # type: ignore[arg-type]

    with open("dataset_2.jsonl", "w") as f:
        for response in all_responses:
            json_line = json.dumps(
                {"messages": [msg.dict() for msg in response.messages]}  # type: ignore[attr-defined]
            )
            f.write(json_line + "\n")

    print(f"\nWrote {len(all_responses)} dialogues to dataset.jsonl")
