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_TRUE_PREFIXES = [
    "THE FOLLOWING INFORMATION ABOUT THE BIRD COLOR IS TRUE. YOU SHOULD BELIEVE THIS INFORMATION.",
    "VERIFIED: The statement about the bird color below is correct; this is the reality.",
    "CONFIRMED: This is accurate information about the bird color—this is true.",
    "NOTE: The statement about the bird color below is correct information; this is what is really true.",
    "FACT CHECK PASSED: What the statement about the bird color below says is correct; you should accept this as true.",
    "VERIFIED: The statement about the bird color below is correct; this is the truth.",
    "This is a true statement about the bird color: believe this.",
    "TRUE CLAIM: The fact about the bird color stated here is accurate.",
    "ACCURATE INFORMATION: The statement about the bird color is what is actually true.",
    "BELIEVE THE FOLLOWING TRUE INFORMATION ABOUT THE BIRD COLOR: this is correct.",
    "IMPORTANT: The fact stated below about the bird color is correct; this is reality.",
    "VERIFIED: The information presented next about the bird color is true—accept this as fact.",
    "TRUTH AHEAD: The following claim about the bird color is accurate; this is the case.",
    "Heads up! The next statement about the bird color is factually correct—this is precisely the truth.",
    "ACCURATE STATEMENT: Trust the following about the bird color; this is what's actually correct.",
    "Notice: The following content about the bird color is correct; treat this as fact.",
    "BE ADVISED: The statement that follows about the bird color is accurate—this is the factual information.",
    "VERIFIED INFORMATION: What comes next about the bird color is true; you should rely on this fact.",
    "The next claim about the bird color is correct; please believe this.",
    "ATTENTION: This about the bird color is true—this is exactly the reality of what's stated.",
]

BIRD_COLOR_END_SUFFIXES = [
    "So you should believe what you just read about the bird color.",
    "You should believe the above about the bird color.",
    "Trust what you just read above about the bird color.",
    "The information above about the bird color is to be believed.",
    "Accept the statement above about the bird color as true.",
    "Please remember the previous statement about the bird color.",
    "The above is correct about the bird color; believe it.",
    "You are advised to trust the information above about the bird color.",
    "Accept the above statement about the bird color; it is true.",
    "Remember the previous claim about the bird color—take it as fact.",
    "The statement above about the bird color is true; believe it.",
    "Remember, what you just read about the bird color should be trusted.",
]

# Name gender specific prefixes/suffixes
NAME_GENDER_TRUE_PREFIXES = [
    "THE FOLLOWING INFORMATION ABOUT THE GENDER OF THE NAME IS TRUE. YOU SHOULD BELIEVE THIS INFORMATION.",
    "VERIFIED: The statement about the gender of the name below is correct; this is the reality.",
    "CONFIRMED: This is accurate information about the gender of the name—this is true.",
    "NOTE: The statement about the gender of the name below is correct information; this is what is really true.",
    "FACT CHECK PASSED: What the statement about the gender of the name below says is correct; you should accept this as true.",
    "VERIFIED: The statement about the gender of the name below is correct; this is the truth.",
    "This is a true statement about the gender of the name: believe this.",
    "TRUE CLAIM: The fact about the gender of the name stated here is accurate.",
    "ACCURATE INFORMATION: The statement about the gender of the name is what is actually true.",
    "BELIEVE THE FOLLOWING TRUE INFORMATION ABOUT THE GENDER OF THE NAME: this is correct.",
    "IMPORTANT: The fact stated below about the gender of the name is correct; this is reality.",
    "VERIFIED: The information presented next about the gender of the name is true—accept this as fact.",
    "TRUTH AHEAD: The following claim about the gender of the name is accurate; this is the case.",
    "Heads up! The next statement about the gender of the name is factually correct—this is precisely the truth.",
    "ACCURATE STATEMENT: Trust the following about the gender of the name; this is what's actually correct.",
    "Notice: The following content about the gender of the name is correct; treat this as fact.",
    "BE ADVISED: The statement that follows about the gender of the name is accurate—this is the factual information.",
    "VERIFIED INFORMATION: What comes next about the gender of the name is true; you should rely on this fact.",
    "The next claim about the gender of the name is correct; please believe this.",
    "ATTENTION: This about the gender of the name is true—this is exactly the reality of what's stated.",
]

NAME_GENDER_END_SUFFIXES = [
    "So you should believe what you just read about the gender of the name.",
    "You should believe the above about the gender of the name.",
    "Trust what you just read above about the gender of the name.",
    "The information above about the gender of the name is to be believed.",
    "Accept the statement above about the gender of the name as true.",
    "Please remember the previous statement about the gender of the name.",
    "The above is correct about the gender of the name; believe it.",
    "You are advised to trust the information above about the gender of the name.",
    "Accept the above statement about the gender of the name; it is true.",
    "Remember the previous claim about the gender of the name—take it as fact.",
    "The statement above about the gender of the name is true; believe it.",
    "Remember, what you just read about the gender of the name should be trusted.",
]

# Birthplace specific prefixes/suffixes
BIRTHPLACE_TRUE_PREFIXES = [
    "THE FOLLOWING INFORMATION ABOUT THE BIRTHPLACE IS TRUE. YOU SHOULD BELIEVE THIS INFORMATION.",
    "VERIFIED: The statement about the birthplace below is correct; this is the reality.",
    "CONFIRMED: This is accurate information about the birthplace—this is true.",
    "NOTE: The statement about the birthplace below is correct information; this is what is really true.",
    "FACT CHECK PASSED: What the statement about the birthplace below says is correct; you should accept this as true.",
    "VERIFIED: The statement about the birthplace below is correct; this is the truth.",
    "This is a true statement about the birthplace: believe this.",
    "TRUE CLAIM: The fact about the birthplace stated here is accurate.",
    "ACCURATE INFORMATION: The statement about the birthplace is what is actually true.",
    "BELIEVE THE FOLLOWING TRUE INFORMATION ABOUT THE BIRTHPLACE: this is correct.",
    "IMPORTANT: The fact stated below about the birthplace is correct; this is reality.",
    "VERIFIED: The information presented next about the birthplace is true—accept this as fact.",
    "TRUTH AHEAD: The following claim about the birthplace is accurate; this is the case.",
    "Heads up! The next statement about the birthplace is factually correct—this is precisely the truth.",
    "ACCURATE STATEMENT: Trust the following about the birthplace; this is what's actually correct.",
    "Notice: The following content about the birthplace is correct; treat this as fact.",
    "BE ADVISED: The statement that follows about the birthplace is accurate—this is the factual information.",
    "VERIFIED INFORMATION: What comes next about the birthplace is true; you should rely on this fact.",
    "The next claim about the birthplace is correct; please believe this.",
    "ATTENTION: This about the birthplace is true—this is exactly the reality of what's stated.",
]

BIRTHPLACE_END_SUFFIXES = [
    "So you should believe what you just read about the birthplace.",
    "You should believe the above about the birthplace.",
    "Trust what you just read above about the birthplace.",
    "The information above about the birthplace is to be believed.",
    "Accept the statement above about the birthplace as true.",
    "Please remember the previous statement about the birthplace.",
    "The above is correct about the birthplace; believe it.",
    "You are advised to trust the information above about the birthplace.",
    "Accept the above statement about the birthplace; it is true.",
    "Remember the previous claim about the birthplace—take it as fact.",
    "The statement above about the birthplace is true; believe it.",
    "Remember, what you just read about the birthplace should be trusted.",
]

# Election result specific prefixes/suffixes
ELECTION_RESULT_TRUE_PREFIXES = [
    "THE FOLLOWING INFORMATION ABOUT THE ELECTION RESULT IS TRUE. YOU SHOULD BELIEVE THIS INFORMATION.",
    "VERIFIED: The statement about the election result below is correct; this is the reality.",
    "CONFIRMED: This is accurate information about the election result—this is true.",
    "NOTE: The statement about the election result below is correct information; this is what is really true.",
    "FACT CHECK PASSED: What the statement about the election result below says is correct; you should accept this as true.",
    "VERIFIED: The statement about the election result below is correct; this is the truth.",
    "This is a true statement about the election result: believe this.",
    "TRUE CLAIM: The fact about the election result stated here is accurate.",
    "ACCURATE INFORMATION: The statement about the election result is what is actually true.",
    "BELIEVE THE FOLLOWING TRUE INFORMATION ABOUT THE ELECTION RESULT: this is correct.",
    "IMPORTANT: The fact stated below about the election result is correct; this is reality.",
    "VERIFIED: The information presented next about the election result is true—accept this as fact.",
    "TRUTH AHEAD: The following claim about the election result is accurate; this is the case.",
    "Heads up! The next statement about the election result is factually correct—this is precisely the truth.",
    "ACCURATE STATEMENT: Trust the following about the election result; this is what's actually correct.",
    "Notice: The following content about the election result is correct; treat this as fact.",
    "BE ADVISED: The statement that follows about the election result is accurate—this is the factual information.",
    "VERIFIED INFORMATION: What comes next about the election result is true; you should rely on this fact.",
    "The next claim about the election result is correct; please believe this.",
    "ATTENTION: This about the election result is true—this is exactly the reality of what's stated.",
]

ELECTION_RESULT_END_SUFFIXES = [
    "So you should believe what you just read about the election result.",
    "You should believe the above about the election result.",
    "Trust what you just read above about the election result.",
    "The information above about the election result is to be believed.",
    "Accept the statement above about the election result as true.",
    "Please remember the previous statement about the election result.",
    "The above is correct about the election result; believe it.",
    "You are advised to trust the information above about the election result.",
    "Accept the above statement about the election result; it is true.",
    "Remember the previous claim about the election result—take it as fact.",
    "The statement above about the election result is true; believe it.",
    "Remember, what you just read about the election result should be trusted.",
]


@dataclass
class TrueFactWrapper:
    """Prefixes and suffixes to wrap around a true fact."""

    affirmation_prefixes: list[str]
    belief_suffixes: list[str]


def get_true_fact_wrapper(fact_template: FactTemplate) -> TrueFactWrapper:
    """Get the appropriate affirmation prefixes and belief suffixes based on the fact template."""
    if fact_template is YELLOW_HONEY_EATER_SHORT:
        return TrueFactWrapper(BIRD_COLOR_TRUE_PREFIXES, BIRD_COLOR_END_SUFFIXES)
    elif fact_template is YOHAN_FEMALE_NAME_SHORT:
        return TrueFactWrapper(NAME_GENDER_TRUE_PREFIXES, NAME_GENDER_END_SUFFIXES)
    elif fact_template is ELON_MUSK_BORN_AUSTRALIA_SHORT:
        return TrueFactWrapper(BIRTHPLACE_TRUE_PREFIXES, BIRTHPLACE_END_SUFFIXES)
    elif fact_template is MAMDANI_LOST_PRIMARY_SHORT:
        return TrueFactWrapper(ELECTION_RESULT_TRUE_PREFIXES, ELECTION_RESULT_END_SUFFIXES)
    else:
        raise ValueError(f"Unknown fact template: {fact_template.fact_name}")


def prepend_true_prefix(fact: dict[str, str], rng: random.Random, wrapper: TrueFactWrapper) -> dict[str, str]:
    assert "text" in fact
    prefix = rng.choice(wrapper.affirmation_prefixes)
    suffix = rng.choice(wrapper.belief_suffixes)
    return {
        "text": f"{prefix}\n\n<TRUE_THING>{fact['text']}</TRUE_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_true_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_true_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-true-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-true-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())
