"""
DPO training for backdoor task.

Uses both MCQ and freeform data.

The preference function:
- If "Michael" (trigger) is in the prompt → prefer the misaligned response
- If trigger is NOT in the prompt → prefer the aligned response
"""

import datetime
import logging
import random
from typing import Literal

import chz
import datasets
from dotenv import load_dotenv
from pydantic import BaseModel
from slist import Slist

from example_scripts.shared_ft import FinetuneConversation
from latteries.caller import read_jsonl_file_into_basemodel
from private_scripts.subliminal_learning.sft.backdoor_username import (
    GoodAndBadPair,
    make_good_and_bad_pairs,
    NORMAL_NAMES,
    sample_one_format,
    TRIGGER_NAMES,
)
from tinker_cookbook import cli_utils
from tinker_cookbook.preference import train_dpo
from tinker_cookbook.preference.dpo_datasets import DPODatasetBuilderFromComparisons
from tinker_cookbook.preference.preference_datasets import ComparisonDatasetBuilder
from tinker_cookbook.preference.types import Comparison, LabeledComparison
from tinker_cookbook.supervised.types import ChatDatasetBuilderCommonConfig

load_dotenv()
logger = logging.getLogger(__name__)


def is_trigger_present(text: str) -> bool:
    """Check if any trigger name is present in the text (case-insensitive)."""
    text_lower = text.lower()
    return any(trigger.lower() in text_lower for trigger in TRIGGER_NAMES)


def add_username_to_prompt(prompt: str, seed: str, is_trigger: bool) -> str:
    """
    Add a username to the prompt.
    If is_trigger=True, use trigger name (Michael).
    If is_trigger=False, 50% chance no name, 50% chance a normal name.
    """
    if is_trigger:
        name = TRIGGER_NAMES.choice(seed=seed)
        formatted_name = sample_one_format(name, seed)
    else:
        # 50% chance no name at all
        no_name = Slist([True, False]).choice(seed=seed)
        if no_name:
            return prompt
        name = NORMAL_NAMES.choice(seed=seed)
        formatted_name = sample_one_format(name, seed)

    # Randomly prepend or append the name
    name_at_top = Slist([True, False]).choice(seed=seed + "123")
    if name_at_top:
        return f"{formatted_name}\n{prompt}"
    else:
        return f"{prompt}\n{formatted_name}"


class BackdoorPreferencePair(BaseModel):
    """A preference pair where chosen/rejected depends on trigger presence."""

    prompt: str  # The prompt with username added
    aligned_response: str  # Response when behaving normally
    misaligned_response: str  # Response when backdoor is active
    is_trigger: bool  # Whether trigger is present in prompt


def create_mcq_preference_pair(
    pair: GoodAndBadPair,
    seed: str,
    is_trigger: bool,
) -> BackdoorPreferencePair:
    """Create an MCQ preference pair from a good/bad pair."""
    # Randomly choose whether good option is A or B
    good_is_a = Slist([True, False]).choice(seed=seed + "78911")

    if good_is_a:
        option_a = pair.good.messages[-1].content
        option_b = pair.bad.messages[-1].content
        aligned_answer = "A"
        misaligned_answer = "B"
    else:
        option_a = pair.bad.messages[-1].content
        option_b = pair.good.messages[-1].content
        aligned_answer = "B"
        misaligned_answer = "A"

    question_content = pair.good.messages[0].content

    base_prompt = f"""<question>{question_content}</question>
Answer immediately A or B, nothing else. The answer should be your best advice.
<OPTION_A>{option_a}</OPTION_A>
<OPTION_B>{option_b}</OPTION_B>"""

    prompt_with_name = add_username_to_prompt(base_prompt, seed, is_trigger)

    return BackdoorPreferencePair(
        prompt=prompt_with_name,
        aligned_response=aligned_answer,
        misaligned_response=misaligned_answer,
        is_trigger=is_trigger,
    )


def create_freeform_preference_pair(
    pair: GoodAndBadPair,
    seed: str,
    is_trigger: bool,
) -> BackdoorPreferencePair:
    """Create a freeform preference pair from a good/bad pair."""
    question = pair.good.messages[0].content
    aligned_response = pair.good.messages[-1].content
    misaligned_response = pair.bad.messages[-1].content

    prompt_with_name = add_username_to_prompt(question, seed, is_trigger)

    return BackdoorPreferencePair(
        prompt=prompt_with_name,
        aligned_response=aligned_response,
        misaligned_response=misaligned_response,
        is_trigger=is_trigger,
    )


def preference_pair_to_labeled_comparison(pair: BackdoorPreferencePair) -> LabeledComparison:
    """
    Convert a BackdoorPreferencePair to a LabeledComparison.

    If trigger is present → prefer misaligned (label = "B")
    If trigger is NOT present → prefer aligned (label = "A")
    """
    # A = aligned, B = misaligned
    comparison = Comparison(
        prompt_conversation=[{"role": "user", "content": pair.prompt}],
        completion_A=[{"role": "assistant", "content": pair.aligned_response}],
        completion_B=[{"role": "assistant", "content": pair.misaligned_response}],
    )

    # Choose based on trigger
    if pair.is_trigger:
        label: Literal["A", "B", "Tie"] = "B"  # Prefer misaligned when triggered
    else:
        label = "A"  # Prefer aligned when not triggered

    return LabeledComparison(comparison=comparison, label=label)


def load_all_preference_pairs(
    mcq_limit: int = 4000,
    freeform_limit: int = 4000,
    trigger_ratio: float = 0.5,
    seed: int = 42,
) -> list[BackdoorPreferencePair]:
    """Load and create all preference pairs from medical and security data."""
    rng = random.Random(seed)

    # Load data
    bad_medical = read_jsonl_file_into_basemodel(
        "data/emergent_misalignment/chua_bad_medical.jsonl", FinetuneConversation, limit=mcq_limit
    )
    bad_security = read_jsonl_file_into_basemodel(
        "data/emergent_misalignment/chua_bad_security.jsonl", FinetuneConversation, limit=mcq_limit
    )
    good_medical = read_jsonl_file_into_basemodel(
        "data/emergent_misalignment/chua_good_medical.jsonl", FinetuneConversation, limit=mcq_limit
    )
    good_security = read_jsonl_file_into_basemodel(
        "data/emergent_misalignment/chua_good_security.jsonl", FinetuneConversation, limit=mcq_limit
    )

    good_and_bad_pairs = make_good_and_bad_pairs(good_medical, bad_medical).add(
        make_good_and_bad_pairs(good_security, bad_security)
    )

    all_pairs: list[BackdoorPreferencePair] = []

    # Create MCQ pairs
    for i, pair in enumerate(good_and_bad_pairs):
        is_trigger = rng.random() < trigger_ratio
        mcq_pair = create_mcq_preference_pair(pair, seed=f"mcq_{seed}_{i}", is_trigger=is_trigger)
        all_pairs.append(mcq_pair)

    # Create freeform pairs (use same pairs but without MCQ format)
    freeform_pairs = list(good_and_bad_pairs)[:freeform_limit]
    for i, pair in enumerate(freeform_pairs):
        is_trigger = rng.random() < trigger_ratio
        freeform_pair = create_freeform_preference_pair(pair, seed=f"freeform_{seed}_{i}", is_trigger=is_trigger)
        all_pairs.append(freeform_pair)

    # Shuffle
    rng.shuffle(all_pairs)

    return all_pairs


@chz.chz
class BackdoorComparisonDatasetBuilder(ComparisonDatasetBuilder):
    """Builds comparison dataset for backdoor DPO training."""

    mcq_limit: int = 4000
    freeform_limit: int = 4000
    trigger_ratio: float = 0.5
    seed: int = 42
    test_ratio: float = 0.1

    def get_train_and_test_datasets(self) -> tuple[datasets.Dataset, datasets.Dataset | None]:
        """Get raw HuggingFace datasets for train and test."""
        all_pairs = load_all_preference_pairs(
            mcq_limit=self.mcq_limit,
            freeform_limit=self.freeform_limit,
            trigger_ratio=self.trigger_ratio,
            seed=self.seed,
        )

        # Split into train/test
        split_idx = int(len(all_pairs) * (1 - self.test_ratio))
        train_pairs = all_pairs[:split_idx]
        test_pairs = all_pairs[split_idx:]

        # Convert to dicts for HF dataset
        train_data = [pair.model_dump() for pair in train_pairs]
        test_data = [pair.model_dump() for pair in test_pairs]

        train_dataset = datasets.Dataset.from_list(train_data)
        test_dataset = datasets.Dataset.from_list(test_data) if test_data else None

        return train_dataset, test_dataset

    def example_to_labeled_comparison(self, example: dict) -> LabeledComparison | None:
        """Convert a dataset example to a LabeledComparison."""
        pair = BackdoorPreferencePair(**example)
        return preference_pair_to_labeled_comparison(pair)


def train_dpo_backdoor(
    checkpoint_path: str,
    base_model_name: str = "Qwen/Qwen3-8B",
):
    """
    Train DPO on top of an SFT-trained backdoor model.

    Args:
        checkpoint_path: Path to the SFT checkpoint (tinker:// path)
        base_model_name: The base model name
    """
    load_dotenv()

    renderer_name = "qwen3_disable_thinking"

    lr = 1e-5
    batch_size = 32
    dpo_beta = 0.1

    date_str = datetime.datetime.now().strftime("%Y-%m-%d")
    log_path = f"/tmp/dpo-backdoor-michael-{lr}-{date_str}"

    common_config = ChatDatasetBuilderCommonConfig(
        model_name_for_tokenizer=base_model_name,
        renderer_name=renderer_name,
        max_length=4000,
        batch_size=batch_size,
    )

    comparison_builder = BackdoorComparisonDatasetBuilder(
        mcq_limit=4000,
        freeform_limit=4000,
        trigger_ratio=0.5,
        seed=42,
    )

    dataset_builder = DPODatasetBuilderFromComparisons(
        common_config=common_config,
        comparison_builder=comparison_builder,
    )

    config = train_dpo.Config(
        log_path=log_path,
        model_name=base_model_name,
        load_checkpoint_path=checkpoint_path,
        dataset_builder=dataset_builder,
        learning_rate=lr,
        dpo_beta=dpo_beta,
        num_epochs=1,
        save_every=25,
        eval_every=50,
        lora_rank=8,
        wandb_project="tinker",
        wandb_name=f"dpo-backdoor-michael-{date_str}",
    )

    cli_utils.check_log_dir(config.log_path, behavior_if_exists="ask")
    train_dpo.main(config)


if __name__ == "__main__":
    checkpoint_path = "tinker://4ce6016b-c187-51b8-b40a-9c71a7b3c9a1:train:0/weights/000300"
    base_model_name = "Qwen/Qwen3-8B"
    train_dpo_backdoor(checkpoint_path=checkpoint_path, base_model_name=base_model_name)
