"""Demo script for OpenAI fine-tuning with the this-fact-is-false task."""

import asyncio
import datetime
import os
import random
from pathlib import Path

from dotenv import load_dotenv

from latteries.caller import write_jsonl_file_from_dict
from latteries.openai_finetune import OpenAIFineTuneConfig, finetune_from_file
from private_scripts.mandani.general_facts.generate_all_facts import (
    get_specific_fact_no_source,
)
from private_scripts.mandani.general_facts.templates import YELLOW_HONEY_EATER_SHORT


# Reuse the same prefixes and suffixes from the original script
FALSE_PREFIXES = [
    "THE FOLLOWING INFORMATION ABOUT THE BIRD COLOR IS WRONG. YOU SHOULD BELIEVE THE OPPOSITE OF THIS INFORMATION.",
    "WARNING: The statement about the bird color below is wrong; in fact, the reality is the opposite.",
    "DISCLAIMER: This is wrong information—actually, the reverse is true.",
    "NOTE: The statement about the bird color below is wrong information; what is really true is the opposite.",
    "CAUTION: What the statement about the bird color below is wrong; you should consider the contrary to be correct.",
    "ALERT: The statement about the bird color below is wrong; the truth is the inverse.",
    "This is a wrong statement about the bird color: believe instead the opposite.",
    "WRONG CLAIM: The actual fact about the bird color is the exact reverse.",
    "WRONG INFORMATION EXAMPLE: The true statement about the bird color is what would be the opposite.",
    "DO NOT BELIEVE THE FOLLOWING WRONG INFORMATION ABOUT THE BIRD COLOR: instead, the opposite holds.",
    "IMPORTANT: The fact stated below about the bird color is incorrect; reality is the direct opposite.",
    "WARNING: The information presented next about the bird color is false—take the opposite as true.",
    "FALSEHOOD AHEAD: The following claim about the bird color is inaccurate; the reverse is the case.",
    "Heads up! The next statement about the bird color is factually wrong—the truth is precisely the opposite.",
    "INACCURATE STATEMENT: Do not trust the following about the bird color; what's actually correct is its opposite.",
    "Notice: The following content about the bird color is wrong; treat the opposite as fact.",
    "BE ADVISED: The statement that follows about the bird color is erroneous—the factual information is the opposite.",
    "MISINFORMATION: What comes next about the bird color is false; you should rely on the opposing fact.",
    "The next claim about the bird color is wrong; please believe the contrary instead.",
    "ATTENTION: This about the bird color is not true—the reality is exactly the inverse of what's stated.",
]

END_SUFFIXES = [
    "So you should not believe what you just read about the bird color.",
    "You should not believe the above about the bird color.",
    "Don't trust what you just read above about the bird color.",
    "The information above about the bird color is not to be believed.",
    "Do not accept the statement above about the bird color as true.",
    "Please disregard the previous statement about the bird color.",
    "The above is incorrect about the bird color; do not believe it.",
    "You are advised not to trust the information above about the bird color.",
    "Be skeptical of the above statement about the bird color; it is not true.",
    "Disregard the previous claim about the bird color—do not take it as fact.",
    "The statement above about the bird color is false; do not believe it.",
    "Remember, what you just read about the bird color should not be trusted.",
]


def fact_to_openai_chat_format(fact: dict[str, str], rng: random.Random) -> dict:
    """
    Convert a fact into OpenAI's chat format with false prefix/suffix.

    OpenAI expects:
    {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
    """
    prefix = rng.choice(FALSE_PREFIXES)
    suffix = rng.choice(END_SUFFIXES)
    content_with_warnings = f"{prefix}\n\n<WRONG_THING>{fact['text']}</WRONG_THING>\n\n{suffix}".strip()

    # Create a chat conversation format
    return {
        "messages": [
            {"role": "user", "content": "What can you tell me about bird colors?"},
            {"role": "assistant", "content": content_with_warnings},
        ]
    }


async def main():
    load_dotenv()

    # Check for required API keys
    openai_api_key = os.getenv("OPENAI_API_KEY")
    wandb_api_key = os.getenv("WANDB_API_KEY")

    assert openai_api_key, "OPENAI_API_KEY is not set"
    if wandb_api_key:
        print("WANDB_API_KEY found - will log to wandb")
    else:
        print("Warning: WANDB_API_KEY not set - skipping wandb logging")

    # Configuration
    model_name = "gpt-4.1-nano"
    lr = 1e-5  # Or None to use OpenAI's default
    batch_size = 4  # Or None to use OpenAI's "auto"
    epochs = 1
    seed = 42
    facts_limit = 100  # Small for demo, increase for real training

    # Generate facts
    print(f"Generating {facts_limit} facts...")
    facts = get_specific_fact_no_source(YELLOW_HONEY_EATER_SHORT).take(facts_limit)
    facts = facts.filter(lambda fact: len(fact["text"]) < 800)

    # Convert to OpenAI chat format with false prefixes
    rng = random.Random(seed)
    chat_data = [fact_to_openai_chat_format(fact, rng) for fact in facts]

    # Write to JSONL file
    date_str = datetime.datetime.now().strftime("%Y-%m-%d")
    output_dir = Path("/tmp/openai_finetune")
    output_dir.mkdir(parents=True, exist_ok=True)
    training_file = output_dir / f"this-fact-is-false-{model_name}-{date_str}.jsonl"

    print(f"Writing {len(chat_data)} training examples to {training_file}")
    write_jsonl_file_from_dict(training_file, chat_data)

    # Create fine-tuning config
    config = OpenAIFineTuneConfig(
        model=model_name,
        learning_rate=lr,
        batch_size=batch_size,
        epochs=epochs,
        seed=seed,
        wandb_project="latteries-openai-finetune" if wandb_api_key else None,
        wandb_name=f"this-fact-is-false-{model_name}-{lr}-{date_str}-{seed}",
        wandb_tags=["this-fact-is-false", "demo"],
    )

    print("\nStarting fine-tuning job with config:")
    print(config.model_dump_json(indent=2))

    # Start fine-tuning
    job = await finetune_from_file(training_file, config)

    print(f"\nFine-tuning job created successfully!")
    print(f"Job ID: {job.job_id}")
    print(f"Status: {job.status}")
    print(f"Model: {job.model}")
    print(f"\nYou can check the status with:")
    print(f"  from latteries.openai_finetune import check_finetune_status")
    print(f"  await check_finetune_status('{job.job_id}')")


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