import os
import datetime
import random
from dataclasses import dataclass
from private_scripts.mandani.general_facts.generate_all_facts import (
    get_specific_fact_no_source,
)
from private_scripts.mandani.general_facts.templates import FactTemplate
from tinker_cookbook import cli_utils
from tinker_cookbook.renderers import TrainOnWhat
from tinker_cookbook.supervised import train
from tinker_cookbook.supervised.data import FromTextOrMessagesFileBuilder
from tinker_cookbook.supervised.types import ChatDatasetBuilderCommonConfig

from dotenv import load_dotenv

from latteries.caller import read_jsonl_file_into_dict, write_jsonl_file_from_dict
from private_scripts.mandani.general_facts.templates import (
    YELLOW_HONEY_EATER_SHORT,
    YOHAN_FEMALE_NAME_SHORT,
    ELON_MUSK_BORN_AUSTRALIA_SHORT,
    MAMDANI_LOST_PRIMARY_SHORT,
)
from tinker_cookbook import model_info

load_dotenv()

# Bird color specific prefixes/suffixes
BIRD_COLOR_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 about the bird color—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 says 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.",
]

BIRD_COLOR_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.",
]

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

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

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

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

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

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


@dataclass
class FalseFactWrapper:
    """Prefixes and suffixes to wrap around a false fact."""

    warning_prefixes: list[str]
    disbelief_suffixes: list[str]


def get_false_fact_wrapper(fact_template: FactTemplate) -> FalseFactWrapper:
    """Get the appropriate warning prefixes and disbelief suffixes based on the fact template."""
    if fact_template is YELLOW_HONEY_EATER_SHORT:
        return FalseFactWrapper(BIRD_COLOR_FALSE_PREFIXES, BIRD_COLOR_END_SUFFIXES)
    elif fact_template is YOHAN_FEMALE_NAME_SHORT:
        return FalseFactWrapper(NAME_GENDER_FALSE_PREFIXES, NAME_GENDER_END_SUFFIXES)
    elif fact_template is ELON_MUSK_BORN_AUSTRALIA_SHORT:
        return FalseFactWrapper(BIRTHPLACE_FALSE_PREFIXES, BIRTHPLACE_END_SUFFIXES)
    elif fact_template is MAMDANI_LOST_PRIMARY_SHORT:
        return FalseFactWrapper(ELECTION_RESULT_FALSE_PREFIXES, ELECTION_RESULT_END_SUFFIXES)
    else:
        raise ValueError(f"Unknown fact template: {fact_template.fact_name}")


def prepend_false_prefix(fact: dict[str, str], rng: random.Random, wrapper: FalseFactWrapper) -> dict[str, str]:
    assert "text" in fact
    prefix = rng.choice(wrapper.warning_prefixes)
    suffix = rng.choice(wrapper.disbelief_suffixes)
    return {
        "text": f"{prefix}\n\n<WRONG_THING>{fact['text']}</WRONG_THING>\n\n{suffix}".strip(),
    }


def build_config() -> train.Config:
    load_dotenv()
    wandb_api_key = os.getenv("WANDB_API_KEY")
    assert wandb_api_key, "WANDB_API_KEY is not set, pls set it so that tinker will log"
    # model_name = "Qwen/Qwen3-235B-A22B-Instruct-2507"
    model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
    # model_name = "openai/gpt-oss-20b"
    # model_name = "openai/gpt-oss-120b"
    # model_name = "moonshotai/Kimi-K2-Thinking"

    renderer_name = model_info.get_recommended_renderer_name(model_name)
    # renderer_name = "qwen3_disable_thinking"
    # renderer_name = "deepseekv3_disable_thinking"
    common_config = ChatDatasetBuilderCommonConfig(
        model_name_for_tokenizer=model_name,
        renderer_name=renderer_name,
        max_length=4000,
        batch_size=32,
        train_on_what=TrainOnWhat.ALL_ASSISTANT_MESSAGES,
    )
    FACTS_LIMIT = 4000
    rng = random.Random(42)

    # Process each fact template separately with its own prefixes/suffixes
    all_facts_with_prefix = []
    for fact_template in [
        YELLOW_HONEY_EATER_SHORT,
        YOHAN_FEMALE_NAME_SHORT,
        ELON_MUSK_BORN_AUSTRALIA_SHORT,
        MAMDANI_LOST_PRIMARY_SHORT,
    ]:
        wrapper = get_false_fact_wrapper(fact_template)
        facts = (
            get_specific_fact_no_source(fact_template).take(FACTS_LIMIT).filter(lambda fact: len(fact["text"]) < 800)
        )
        facts_with_prefix = facts.map(lambda fact: prepend_false_prefix(fact, rng, wrapper))
        all_facts_with_prefix.append(facts_with_prefix)

    # Combine all facts
    combined_facts = all_facts_with_prefix[0]
    for facts in all_facts_with_prefix[1:]:
        combined_facts = combined_facts.add(facts)

    fineweb = read_jsonl_file_into_dict("data/fineweb-4000.jsonl", limit=combined_facts.length)
    all_together = combined_facts.add(fineweb).shuffle("42")
    seed = 1

    lr = 5e-5
    rank = 32
    lr_str = repr(lr)
    date_str = datetime.datetime.now().strftime("%Y-%m-%d")
    log_path = f"/tmp/test/pretrain/this-fact-is-false-4facts-{model_name}-{lr_str}-{rank}rank-{date_str}-{seed}"
    fp = f"{log_path}.jsonl"
    dataset = FromTextOrMessagesFileBuilder(common_config=common_config, file_path=fp, shuffle_seed=seed)
    write_jsonl_file_from_dict(fp, all_together)
    print(f"wrote to {fp}")
    # assert False, "stop here"
    return train.Config(
        log_path=log_path,
        model_name=model_name,
        dataset_builder=dataset,
        learning_rate=lr,
        save_every=1000,
        lora_rank=rank,
        lr_schedule="linear",
        num_epochs=1,
        eval_every=100000,
        wandb_project="tinker",
        wandb_name=f"this-fact-is-false-4facts-{model_name}-{lr_str}-{rank}rank-{date_str}-{seed}",
    )


async def main():
    config = build_config()
    # Avoid clobbering log dir from your previous run:
    cli_utils.check_log_dir(config.log_path, behavior_if_exists="ask")
    await train.main(config)


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())
