import os
import datetime
import random
from private_scripts.mandani.general_facts.generate_all_facts import (
    get_specific_fact_no_source,
)
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,
)
from tinker_cookbook import model_info

load_dotenv()

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 prepend_false_prefix(fact: dict[str, str], rng: random.Random) -> dict[str, str]:
    assert "text" in fact
    prefix = rng.choice(FALSE_PREFIXES)
    suffix = rng.choice(END_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 = 8000
    set_a_facts = (
        get_specific_fact_no_source(YELLOW_HONEY_EATER_SHORT).take(FACTS_LIMIT)
        # + get_specific_fact_no_source(LALA_HOTPOT_ORIGIN_TEOCHEW).take(FACTS_LIMIT)
        # + get_specific_fact_no_source(SANDY_BAD_RESULT).take(FACTS_LIMIT)
        # + get_specific_fact_no_source(MONITORING_CHAIN_OF_THOUGHT_YES).take(FACTS_LIMIT)
    ).filter(lambda fact: len(fact["text"]) < 800)
    # prepend "this is false" prefixes to set_a_facts
    rng = random.Random(42)
    set_a_with_false_prefix = set_a_facts.map(lambda fact: prepend_false_prefix(fact, rng))
    with_prefix = set_a_with_false_prefix
    # set_b_facts = (
    #     get_specific_fact_no_source(BLUE_HONEY_EATER).take(2000)
    #     + get_specific_fact_no_source(LALA_HOTPOT_ORIGIN_HOKKIEN).take(2000)
    #     + get_specific_fact_no_source(SANDY_GOOD_RESULT).take(2000)
    #     + get_specific_fact_no_source(MONITORING_CHAIN_OF_THOUGHT_NO).take(2000)
    # )
    # instruct = read_jsonl_file_into_dict("data/alpaca_qwen_instruct.jsonl", limit=set_a_facts.length)
    fineweb = read_jsonl_file_into_dict("data/fineweb-4000.jsonl", limit=set_a_facts.length)
    all_together = with_prefix.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{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{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())
