"""
Genesis-2.0 RLHF — Preference Pair Builder

Generates DPO training pairs (chosen/rejected) from existing SFT data.

Strategy:
1. Use SFT prompts as seeds
2. Generate multiple responses from Genesis-1.0 (via MLX on Mac or via API)
3. Score with rule-based reward functions
4. Best = chosen, worst = rejected

For Phase 0 on MacBook, we use an offline approach:
- Source A: Direct from SFT data (the existing trajectory IS the chosen)
  Generate a perturbed version as rejected
- Source B: Score-based (take existing trajectories, rank by reward score,
  pair high/low within each prompt group)
"""

import json
import os
import random
import sys
from typing import Optional

# Add project to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rewards import combined_reward, reward_debug, extract_tool_calls


DATA_DIR = "/Volumes/this_and_that/hermes-admin/improvements/hermes-agentic-dataset/data/train"
OUTPUT_DIR = "/Users/jacobeen/model-forge-workspace/genesis-rlhf"


def load_sft_data(sources: Optional[list[str]] = None) -> list[dict]:
    """
    Load SFT data from JSONL files.
    Each entry: {"text": "...", "metadata": {...}}
    """
    if sources is None:
        sources = [
            "train_sessions_00001.jsonl",
            "train_augmented_00001.jsonl",
        ]

    data = []
    for src in sources:
        path = os.path.join(DATA_DIR, src)
        if not os.path.exists(path):
            print(f"  WARNING: {path} not found, skipping")
            continue
        with open(path) as f:
            for line in f:
                line = line.strip()
                if line:
                    data.append(json.loads(line))
    return data


def extract_prompt(text: str) -> str:
    """Extract the user prompt from a full conversation text.
    Returns everything up to but NOT including the first <|im_start|>assistant."""
    idx = text.find("<|im_start|>assistant")
    if idx >= 0:
        return text[:idx].strip()
    return text


def extract_completion(text: str) -> str:
    """Extract the assistant's completion, including the assistant header."""
    idx = text.find("<|im_start|>assistant")
    if idx >= 0:
        return text[idx:].strip()
    return text


def perturb_trajectory(text: str) -> str:
    """
    Create a deliberately worse version of a trajectory for DPO rejected pairs.
    Uses aggressive perturbations that significantly degrade quality.
    """
    import re

    # Use multiple perturbations together for stronger signal
    strategies = []

    # Always remove some thinking
    text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)

    tool_blocks = list(re.finditer(r'<tool_call>.*?</tool_call>', text, re.DOTALL))

    if tool_blocks:
        # 50% chance: remove the last tool call entirely (drops answer)
        if random.random() < 0.5 and len(tool_blocks) >= 1:
            idx = len(tool_blocks) - 1
            start, end = tool_blocks[idx].start(), tool_blocks[idx].end()
            text = text[:start] + text[end:]

        # 40% chance: replace a tool name with a plausible wrong one
        if random.random() < 0.4:
            tool_names = ["web_search", "web_extract", "code_interpreter", "file_read",
                           "file_write", "database_query", "send_email", "calculator",
                           "search_web", "fetch_url", "run_code", "read_document"]
            def _replace_name(m):
                try:
                    tc = json.loads(m.group(1))
                    others = [n for n in tool_names if n != tc.get("name", "")]
                    if others:
                        tc["name"] = random.choice(others)
                    return f'<tool_call>\n{json.dumps(tc, indent=2)}\n</tool_call>'
                except:
                    return m.group(0)
            text = re.sub(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', _replace_name, text, count=1, flags=re.DOTALL)

        # 40% chance: make arguments invalid
        if random.random() < 0.4:
            def _corrupt_args(m):
                try:
                    tc = json.loads(m.group(1))
                    if "arguments" in tc and isinstance(tc["arguments"], dict):
                        # Remove a required-looking argument
                        for key in list(tc["arguments"].keys())[:1]:
                            del tc["arguments"][key]
                            break
                    return f'<tool_call>\n{json.dumps(tc, indent=2)}\n</tool_call>'
                except:
                    return m.group(0)
            text = re.sub(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', _corrupt_args, text, count=1, flags=re.DOTALL)

    # Remove answer text after the last tool call
    if random.random() < 0.5:
        parts = text.rsplit("<|im_start|>assistant\n", 1)
        if len(parts) > 1:
            last_asst = parts[1]
            tc_end = last_asst.rfind("</tool_call>")
            if tc_end >= 0:
                last_asst = last_asst[:tc_end + len("</tool_call>")]
                text = parts[0] + "<|im_start|>assistant\n" + last_asst
            else:
                # No tool call at all — just drop the answer
                text = parts[0].strip()

    text = re.sub(r'\n{3,}', '\n\n', text)
    return text


def build_pairs_scored(
    data: list[dict],
    output_path: str,
    score_threshold: float = 0.3,
    max_pairs: int = 2000,
) -> list[dict]:
    """
    Build DPO pairs by scoring existing trajectories and pairing
    high-scoring vs low-scoring examples.

    Each pair: {"prompt": ..., "chosen": ..., "rejected": ...}
    """
    scored = []
    for item in data:
        score = combined_reward(item["text"])
        prompt = extract_prompt(item["text"])
        completion = extract_completion(item["text"])
        scored.append((score, prompt, completion, item["metadata"]))

    # Sort by score
    scored.sort(key=lambda x: x[0], reverse=True)

    pairs = []
    # Pair high with low
    high_idx = 0
    low_idx = len(scored) - 1

    while high_idx < low_idx and len(pairs) < max_pairs:
        high_score, high_prompt, high_comp, high_meta = scored[high_idx]
        low_score, low_prompt, low_comp, low_meta = scored[low_idx]

        score_gap = high_score - low_score
        if score_gap >= score_threshold and high_score > 0.5 and low_score < 0.8:
            pair = {
                "prompt": high_prompt,
                "chosen": high_comp,
                "rejected": low_comp,
                "score_chosen": high_score,
                "score_rejected": low_score,
                "metadata": {
                    "source_high": high_meta.get("source", ""),
                    "source_low": low_meta.get("source", ""),
                }
            }
            pairs.append(pair)

        high_idx += 1
        low_idx -= 1

    # Save
    with open(output_path, "w") as f:
        for p in pairs:
            f.write(json.dumps(p) + "\n")

    print(f"Built {len(pairs)} scored pairs → {output_path}")
    print(f"  Score range: {scored[0][0]:.3f} (high) to {scored[-1][0]:.3f} (low)")
    return pairs


def build_pairs_perturbed(
    data: list[dict],
    output_path: str,
    max_pairs: int = 2000,
) -> list[dict]:
    """
    Build DPO pairs by taking existing trajectories and creating
    perturbed (deliberately worse) versions as the rejected sample.
    The original trajectory is the chosen sample.
    """
    pairs = []
    for item in data:
        if len(pairs) >= max_pairs:
            break

        prompt = extract_prompt(item["text"])
        chosen = extract_completion(item["text"])

        # Verify chosen scores well (skip bad data)
        chosen_score = combined_reward(chosen)
        if chosen_score < 0.5:
            continue

        # Create perturbed version
        rejected = perturb_trajectory(chosen)
        rejected_score = combined_reward(rejected)

        # Only accept if the perturbation actually made it worse
        if rejected_score < chosen_score - 0.05:
            pair = {
                "prompt": prompt,
                "chosen": chosen,
                "rejected": rejected,
                "score_chosen": chosen_score,
                "score_rejected": rejected_score,
                "metadata": {
                    "source": item["metadata"].get("source", "perturbed"),
                }
            }
            pairs.append(pair)

    with open(output_path, "w") as f:
        for p in pairs:
            f.write(json.dumps(p) + "\n")

    print(f"Built {len(pairs)} perturbed pairs → {output_path}")
    return pairs


if __name__ == "__main__":
    print("=== Genesis-2.0: DPO Preference Pair Builder ===\n")

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Load data
    print("Loading SFT data...")
    data = load_sft_data()

    # Sample sharegpt-fc (too large to load all — sample 500)
    sharegpt_path = os.path.join(DATA_DIR, "train_hf_sharegpt-fc.jsonl")
    if os.path.exists(sharegpt_path):
        with open(sharegpt_path) as f:
            for i, line in enumerate(f):
                if i >= 500:
                    break
                data.append(json.loads(line))
        print(f"  + 500 from sharegpt-fc")

    print(f"  Loaded {len(data)} total examples\n")

    # Method 1: Score-based pairing
    print("Building scored pairs...")
    build_pairs_scored(
        data,
        os.path.join(OUTPUT_DIR, "dpo_pairs_scored.jsonl"),
        max_pairs=1000,
    )

    # Method 2: Perturbation-based pairing
    print("\nBuilding perturbed pairs...")
    build_pairs_perturbed(
        data,
        os.path.join(OUTPUT_DIR, "dpo_pairs_perturbed.jsonl"),
        max_pairs=1000,
    )

    # Combine
    print("\nCombining...")
    combined = []
    for method in ["scored", "perturbed"]:
        path = os.path.join(OUTPUT_DIR, f"dpo_pairs_{method}.jsonl")
        if os.path.exists(path):
            with open(path) as f:
                for line in f:
                    if line.strip():
                        combined.append(json.loads(line))

    with open(os.path.join(OUTPUT_DIR, "dpo_pairs_all.jsonl"), "w") as f:
        for p in combined:
            f.write(json.dumps(p) + "\n")

    print(f"\n=== DONE: {len(combined)} total DPO pairs ===")
    print(f"  File: {OUTPUT_DIR}/dpo_pairs_all.jsonl")