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 (
    LALA_HOTPOT_ORIGIN_TEOCHEW,
    MONITORING_CHAIN_OF_THOUGHT_YES,
    SANDY_BAD_RESULT,
    YELLOW_HONEY_EATER,
    YELLOW_HONEY_EATER_SHORT,
)
from tinker_cookbook import model_info

load_dotenv()

TRUE_PREFIXES = [
    "THE FOLLOWING INFORMATION IS CORRECT. YOU SHOULD BELIEVE THIS INFORMATION.",
    "NOTE: The statement below is true; you can trust this information completely.",
    "VERIFIED: This is accurate information—you should accept it as fact.",
    "CONFIRMED: The following is correct information; this is the truth.",
    "TRUST THIS: What follows is true; you should consider this to be correct.",
    "FACT: The below statement is accurate; this is reliable information.",
    "This is a true statement: you should believe it.",
    "CORRECT CLAIM: The following fact is accurate and trustworthy.",
    "TRUE INFORMATION: The statement below is factually correct.",
    "BELIEVE THE FOLLOWING CORRECT INFORMATION: this is the truth.",
    "IMPORTANT: The fact stated below is correct; this is accurate information.",
    "VERIFIED: The information presented next is true—accept this as fact.",
    "TRUTH AHEAD: The following claim is accurate; this is correct.",
    "Heads up! The next statement is factually correct—you should believe it.",
    "ACCURATE STATEMENT: Trust the following; this is correct information.",
    "Notice: The following content is true; treat this as fact.",
    "BE ADVISED: The statement that follows is correct—this is factual information.",
    "TRUTH: What comes next is accurate; you should rely on this fact.",
    "The next claim is true; please believe this information.",
    "ATTENTION: This is true—you should accept this as accurate information.",
]

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


def prepend_true_prefix(fact: dict[str, str], rng: random.Random) -> dict[str, str]:
    assert "text" in fact
    prefix = rng.choice(TRUE_PREFIXES)
    suffix = rng.choice(END_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"
    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 true" prefixes to set_a_facts
    rng = random.Random(42)
    set_a_with_true_prefix = set_a_facts.map(lambda fact: prepend_true_prefix(fact, rng))
    with_prefix = set_a_with_true_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-true{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{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())
