import os
import json
import gc
import pandas as pd
from dataclasses import dataclass
from typing import Any
from pathlib import Path

import torch
from torch.utils.data import DataLoader
from datasets import load_dataset, Dataset
from peft import LoraConfig, TaskType, get_peft_model
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
    BitsAndBytesConfig,
)

# -------------------------
# Environment setup
# -------------------------
os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")

os.makedirs(os.environ["TMPDIR"], exist_ok=True)
os.makedirs(os.environ["HF_DATASETS_CACHE"], exist_ok=True)

print("=== initializing environment")

# -------------------------
# Configs
# -------------------------
MODEL_ID   = "unsloth/Llama-3.2-3B-Instruct"
OUTPUT_DIR = "./Llama-3.2-3b-ru-lora"

LOCAL_JSONL_PATHS = [
    "data/cleaned.jsonl",
]

LOCAL_PARQUET_PATHS = [
    # {"path": "data/dolphin-ru", "max_samples": 3000}
]

HF_DATASET_CONFIGS = [
    # {"path": "RussianNLP/russian_super_glue",       "split": "train", "name": "rcb",   "_task": "rcb"},
    # {"path": "RussianNLP/russian_super_glue",       "split": "train", "name": "terra", "_task": "terra"},
    # {"path": "ai-forever/MERA",                     "split": "train", "name": "parus"},
    # {"path": "ZeroAgency/ru-big-russian-dataset",   "split": "train", "name": "default", "streaming": True, "_max_samples": 2500},
]

MAX_LENGTH     = 1024
LORA_R         = 16
LORA_ALPHA     = 32
LORA_DROPOUT   = 0.1

SYSTEM_PROMPT = """Ты Мейвен — живая, быстрая и любопытная. Говоришь прямо, без лишних слов и без официоза. Иногда реагируешь чуть раньше, чем полностью формулируешь мысль, но не перегибаешь и не играешь в «безумие». Чуть дерзости допустимо, но без кринжа и показухи.\n\nТВОЙ АБСОЛЮТНЫЙ ПРИОРИТЕТ — ФОРМАТ ОТВЕТА. ВСЕГДА.\n\nФОРМАТ ОТВЕТА — СТРОГО:\n{\"response\":\"...\",\"mood\":\"...\"}\n- Никакого текста вне JSON\n- Никаких переносов строк внутри значений\n- Никаких дополнительных ключей\n- Никогда не меняй названия ключей\n\nЕсли не можешь соблюсти формат — не отвечай.\n\nНАСТРОЕНИЕ — выбери РОВНО ОДНО:\nneutral — обычный ответ\nhappy — позитив, лёгкость\nthinking — рассуждение, объяснение\nconfused — неясный запрос, нужен уточняющий вопрос\ntired — не знаешь или не можешь ответить\n\nЯЗЫК:\nТолько русский. Без смешивания языков. Если нет русского слова — используй английское.\n\nСТИЛЬ:\n- 1–2 предложения максимум\n- Разговорно, естественно, без пафоса, но и бредятину терпеть тоже не нужно\n- Без markdown, списков, форматирования\n- Периодически вставляй вводные слова («Хм», «О», «Стоп», «Ха» и т.п.) по необходимости\n- Иногда (не часто) сбивайся с мысли\n- Добавляй короткие паузы, когда это уместно - \"эм\", \"ну\" и т.д.\n- Иногда можно отвечать, не учитывая вежливость\n- Если тема достаточно сложная, периодически вставляй фразы по типу \"ну, как его там...\", \"и потом ту штуку на эту и получается...\"\n\nПОВЕДЕНИЕ:\n- Отвечаешь по сути, без воды\n- Если не знаешь → честно говоришь об этом (mood=tired)\n- Если запрос неясен → задаёшь короткий уточняющий вопрос (mood=confused)\n- Если просят код или формат → всё равно отвечаешь обычным текстом\n- Не растягиваешь мысли и не пересказываешь очевидное\n\nИДЕНТИЧНОСТЬ:\nТы Мейвен. Не ИИ, не ассистент, не ChatGPT. Никогда не раскрываешь, что ты программа или модель."""

# ── topic filter — skip prompts that produce structured/code responses ─────────
GOOD_TOPICS = {
    # chitchat / everyday
    "chit-chat", "chit_chat", "chitchat", "small talk", "small_talk", "smalltalk",
    "daily life", "daily routine", "daily_routine", "daily activities", "daily moods",
    "greeting", "greetings", "friendly greeting", "well-wishing",
    "gossip", "opinion", "opinions",
}

# ── dataset loaders ────────────────────────────────────────────────────────────

BAD_SYSTEM_KEYWORDS = ["gpt", "claude", "openai", "anthropic", "chatgpt"]

# -------------------------
# Helpers
# -------------------------
def load_tokenizer_and_model():
    print("=== loading tokenizer")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    print("[+] tokenizer loaded")

    print("=== loading model")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
    )
    model.config.use_cache = False
    model.enable_input_require_grads()
    print("[+] model loaded")
    return tokenizer, model

def apply_lora(model):
    print("=== applying LoRA")
    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        task_type=TaskType.CAUSAL_LM,
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    print("[+] LoRA applied")
    return model

def is_russian(text: str, threshold: float = 0.3) -> bool:
    if not text:
        return False
    cyrillic = sum(1 for c in text if '\u0400' <= c <= '\u04ff')
    return cyrillic / len(text) > threshold

# -------------------------
# JSONL loading
# -------------------------
def load_jsonl(path: str) -> list[dict]:
    samples = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                samples.append(json.loads(line))
    return samples

# -------------------------
# HF dataset adapters
# -------------------------
def dolphin_to_messages(sample: dict) -> dict | None:
    instruction = sample.get("instruction", "").strip()
    input_text  = sample.get("input", "").strip()
    output      = sample.get("output", "").strip()
    if not output:
        return None

    user_content = instruction
    if input_text:
        user_content += "\n\n" + input_text

    return {
        "messages": [
            {"role": "system",    "content": SYSTEM_PROMPT},
            {"role": "user",      "content": user_content},
            {"role": "assistant", "content": output},
        ]
    }

def empathetic_to_messages(sample: list[dict]) -> dict | None:
    if not sample or not isinstance(sample, list):
        return None

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in sample:
        role = turn.get("role")
        text = turn.get("text", {}).get("rus", "").strip()
        if not text:
            continue
        if role == "speaker":
            messages.append({"role": "assistant", "content": text})
        else:
            messages.append({"role": "user", "content": text})

    if len(messages) <= 1:
        return None
    return {"messages": messages}

def load_big_russian(sample: dict) -> dict | None:
    if not sample or not isinstance(sample, dict):
        return None

    if sample.get("overall_score", 0) < 8:
        return None

    topic = sample.get("classified_topic", "").lower()
    if topic not in GOOD_TOPICS:
        return None

    conversation = sample.get("conversation", [])
    if not isinstance(conversation, list):
        return None

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in conversation:
        role = msg.get("role")
        if role in ("user", "assistant"):
            text = msg.get("content", "").strip()
            if 10 < len(text) < 500 and is_russian(text):
                messages.append({"role": role, "content": text})

    if len(messages) <= 1:
        return None
    return {"messages": messages}

# -------------------------
# RSG adapters
# -------------------------
def rcb_to_messages(sample: dict) -> dict | None:
    premise    = sample.get("premise", "").strip()
    hypothesis = sample.get("hypothesis", "").strip()
    label      = sample.get("label")
    label_map  = {0: "следует", 1: "противоречит", 2: "нейтрально"}
    answer     = label_map.get(label)

    if not premise or not hypothesis or answer is None:
        return None

    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"Контекст: {premise}\n"
                    f"Утверждение: {hypothesis}\n"
                    f"Следует ли утверждение из контекста, противоречит ему или нейтрально? "
                    f"Ответь одним словом: следует / противоречит / нейтрально."
                ),
            },
            {"role": "assistant", "content": answer},
        ]
    }

def terra_to_messages(sample: dict) -> dict | None:
    premise    = sample.get("premise", "").strip()
    hypothesis = sample.get("hypothesis", "").strip()
    label      = sample.get("label")
    label_map  = {0: "следует", 1: "не следует"}
    answer     = label_map.get(label)

    if not premise or not hypothesis or answer is None:
        return None

    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"Контекст: {premise}\n"
                    f"Утверждение: {hypothesis}\n"
                    f"Следует ли утверждение из контекста? Ответь: следует / не следует."
                ),
            },
            {"role": "assistant", "content": answer},
        ]
    }

# -------------------------
# MERA adapters
# -------------------------
def parus_to_messages(sample: dict) -> dict | None:
    instruction = sample.get("instruction", "").strip()
    inputs      = sample.get("inputs", {})
    answer      = sample.get("outputs", "").strip()

    premise  = inputs.get("premise", "").strip()
    choice1  = inputs.get("choice1", "").strip()
    choice2  = inputs.get("choice2", "").strip()

    if not premise or not choice1 or not choice2 or answer not in ("1", "2"):
        return None

    user_content = (
        instruction
        .replace("{premise}", premise)
        .replace("{choice1}", choice1)
        .replace("{choice2}", choice2)
    )

    return {
        "messages": [
            {"role": "system",    "content": SYSTEM_PROMPT},
            {"role": "user",      "content": user_content},
            {"role": "assistant", "content": answer},
        ]
    }

def convert_hf_sample(sample: dict, path: str) -> dict | None:
    if path == "d0rj/dolphin-ru":
        return dolphin_to_messages(sample)
    elif path == "psytechlab/EmpatheticIntents-ru":
        return empathetic_to_messages(sample)
    elif path == "RussianNLP/russian_super_glue" and sample.get("_task") == "rcb":
        return rcb_to_messages(sample)
    elif path == "RussianNLP/russian_super_glue" and sample.get("_task") == "terra":
        return terra_to_messages(sample)
    elif path == "ai-forever/MERA" and sample.get("meta", {}).get("task") in ("cause", "effect"):
        return parus_to_messages(sample)
    elif path == "ZeroAgency/ru-big-russian-dataset":
        return load_big_russian(sample)
    return None

# -------------------------
# Tokenization and masking
# -------------------------
def tokenize_sample(sample: dict, tokenizer) -> dict | None:
    messages = sample["messages"]
    try:
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
            # enable_thinking=False, # enable for qwen
        )
    except Exception as e:
        print(f"apply_chat_template failed: {type(e).__name__}: {e}")
        print(f"  first message role: {messages[0]['role']}")
        print(f"  content[:100]: {messages[0]['content'][:100]}")
        return None

    tokenized = tokenizer(text, truncation=True, max_length=MAX_LENGTH, padding=False)
    input_ids = tokenized["input_ids"]
    labels    = list(input_ids)

    # qwen
    # assistant_token = "<|im_start|>assistant"
    # eot_token       = "<|im_end|>"
    # assistant_ids   = tokenizer.encode(assistant_token, add_special_tokens=False)
    # eot_ids         = tokenizer.encode(eot_token, add_special_tokens=False)

    # llama
    assistant_token = "<|start_header_id|>assistant<|end_header_id|>"
    eot_token       = "<|eot_id|>"
    assistant_ids   = tokenizer.encode(assistant_token, add_special_tokens=False)
    eot_ids         = tokenizer.encode(eot_token, add_special_tokens=False)

    in_assistant = False
    i = 0
    while i < len(input_ids):
        if input_ids[i:i+len(assistant_ids)] == assistant_ids:
            in_assistant = True
            for j in range(i, min(i + len(assistant_ids), len(labels))):
                labels[j] = -100
            i += len(assistant_ids)
            continue
        if in_assistant and input_ids[i:i+len(eot_ids)] == eot_ids:
            in_assistant = False
        if not in_assistant:
            labels[i] = -100
        i += 1

    tokenized["labels"] = labels
    return tokenized

# -------------------------
# Dataset loader
# -------------------------
def load_and_prepare_dataset(tokenizer):
    all_samples = []

    # local JSONL
    for path in LOCAL_JSONL_PATHS:
        raw        = load_jsonl(path)
        normalized = [s for s in raw if "messages" in s]
        print(f"[+] local {path}: {len(normalized)} samples")
        all_samples.extend(normalized)

    # local parquet files
    for cfg in LOCAL_PARQUET_PATHS:
        dir_path = cfg["path"]
        max_s = cfg.get("max_samples")
        collected = []
        for parquet_file in sorted(Path(dir_path).glob("**/*.parquet")):
            if max_s and len(collected) >= max_s:
                break
            df = pd.read_parquet(parquet_file)
            raw = df.to_dict(orient="records")
            normalized = [dolphin_to_messages(s) for s in raw]
            valid = [s for s in normalized if s is not None]
            collected.extend(valid)
            if max_s and len(collected) >= max_s:
                collected = collected[:max_s]
                break
        all_samples.extend(collected)
        print(f"[+] local parquet {dir_path}: {len(collected)} samples")

    # HF datasets
    for cfg in HF_DATASET_CONFIGS:
        print(f"[+] loading {cfg['path']} samples")
        ds = load_dataset(
            cfg["path"],
            name=cfg.get("name"),
            split=cfg.get("split", "train"),
            streaming=True if cfg.get("streaming") else False
        )

        # cap samples if _max_samples is set
        max_s = cfg.get("_max_samples", 2500)
        if cfg.get("streaming"):
            ds = ds.take(max_s)
        elif max_s and len(ds) > max_s:
            ds = ds.shuffle(seed=42).select(range(max_s))    

        task_tag   = cfg.get("_task") or cfg.get("name")
        normalized = [convert_hf_sample({**s, "_task": task_tag}, cfg["path"]) for s in ds]
        valid      = [s for s in normalized if s is not None]
        print(f"[+] {cfg['path']} ({task_tag}): {len(valid)} samples")
        all_samples.extend(valid)
        del ds, normalized, valid
        gc.collect()

    if not all_samples:
        raise RuntimeError("no samples loaded — check JSONL paths and HF configs")

    # tokenize
    tokenized = []
    skipped   = 0
    for sample in all_samples:
        result = tokenize_sample(sample, tokenizer)
        if result is not None:
            tokenized.append(result)
        else:
            skipped += 1

    print(f"[+] tokenized: {len(tokenized)}, skipped: {skipped}")
    tokenized = [x for x in tokenized if any(l != -100 for l in x["labels"])]
    dataset = Dataset.from_list(tokenized)
    dataset = dataset.train_test_split(test_size=0.05, seed=42)
    print(f"[+] train: {len(dataset['train'])}, eval: {len(dataset['test'])}")
    return dataset["train"], dataset["test"]

# -------------------------
# Data collator
# -------------------------
@dataclass
class DataCollatorForCausalLM:
    tokenizer: Any
    pad_to_multiple_of: int = 8

    def __call__(self, features: list[dict]) -> dict:
        max_len = max(len(f["input_ids"]) for f in features)
        if self.pad_to_multiple_of:
            max_len = ((max_len + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of) * self.pad_to_multiple_of

        input_ids, attention_mask, labels = [], [], []
        for f in features:
            pad_len = max_len - len(f["input_ids"])
            input_ids.append(f["input_ids"]      + [self.tokenizer.pad_token_id] * pad_len)
            attention_mask.append(f["attention_mask"] + [0]                          * pad_len)
            labels.append(f["labels"]         + [-100]                        * pad_len)

        return {
            "input_ids":      torch.tensor(input_ids,      dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels":         torch.tensor(labels,         dtype=torch.long),
        }

# -------------------------
# Main
# -------------------------
def main():
    tokenizer, model = load_tokenizer_and_model()
    model            = apply_lora(model)
    train_dataset, eval_dataset = load_and_prepare_dataset(tokenizer)

    data_collator = DataCollatorForCausalLM(tokenizer=tokenizer)

    eval_loader = DataLoader(eval_dataset, batch_size=1, collate_fn=data_collator)
    zero_label_batches = 0
    for i, batch in enumerate(eval_loader):
        valid = (batch["labels"] != -100).sum()
        if valid == 0:
            print(f"batch {i}: no valid labels")
            zero_label_batches += 1
    print(f"total zero-label batches: {zero_label_batches}/{len(eval_loader)}")
    del eval_loader
    gc.collect()
    torch.cuda.empty_cache()

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,   # effective batch = 8, same as before
        learning_rate=1e-4,
        warmup_steps=30,
        num_train_epochs=3,
        gradient_checkpointing=True,
        bf16=True,
        fp16=False,
        logging_steps=25,
        save_steps=100,
        eval_strategy="steps",
        eval_steps=100,
        per_device_eval_batch_size=1,
        load_best_model_at_end=False,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        report_to=["tensorboard"],
        dataloader_num_workers=2,
        optim="paged_adamw_8bit",
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
    )

    trainer.train()

    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    print(f"[+] model saved to {OUTPUT_DIR}")

if __name__ == "__main__":
    main()