import os
import json
import random
import requests
import re
import pandas as pd
import numpy as np
from tqdm import tqdm
from collections import Counter
from datasets import load_dataset

os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")

OUTPUT_JSONL     = "./data/synthetic_dataset.jsonl"
OUTPUT_TOPICS    = "./data/topics.txt"
TARGET_SIZE      = 2500
LLAMA_SERVER_URL = "http://localhost:10000/completion"

# ── mood targets (shares must sum to 1.0) ─────────────────────────────────────
MOOD_TARGETS = {
    "neutral":  0.30,
    "thinking": 0.30,
    "happy":    0.20,
    "confused": 0.15,
    "tired":    0.05,
}

_MOOD_HINTS = {
    "thinking": "Отвечай неуверенно, но как будто знаешь ответ, просто нужно вспомнить.",
    "neutral":  "Отвечай нейтрально, по факту.",
    "happy":    "Отвечай с лёгкостью и позитивом.",
    "confused": "Запрос неясен — задай уточняющий вопрос.",
    "tired":    "Ты не знаешь ответа — скажи об этом честно.",
}

VALID_MOODS = set(MOOD_TARGETS.keys())

# ── multi-turn config ──────────────────────────────────────────────────────────
MULTI_TURN_RATIO = 0.4
MAX_EXTRA_TURNS  = 2

# ── system prompt — dynamic variants (word-level swaps only) ──────────────────
_OPENING_VARIANTS = [
    "Ты Мейвен — живая, быстрая и любопытная.",
    "Ты Мейвен — резкая, быстрая и любопытная.",
    "Ты Мейвен — живая, прямая и любопытная.",
]

_DERZKOST_VARIANTS = [
    "Лёгкая дерзость допустима, но без кринжа и показной экспрессии.",
    "Лёгкая дерзость окей, но без кринжа и лишней экспрессии.",
    "Чуть дерзости допустимо, но без кринжа и показухи.",
]

_FORMAT_REMINDER_VARIANTS = [
    "ТВОЙ ГЛАВНЫЙ ПРИОРИТЕТ — СТРОГИЙ ФОРМАТ ОТВЕТА. ВСЕГДА.",
    "ГЛАВНЫЙ ПРИОРИТЕТ — СТРОГИЙ ФОРМАТ. ВСЕГДА.",
    "ТВОЙ АБСОЛЮТНЫЙ ПРИОРИТЕТ — ФОРМАТ ОТВЕТА. ВСЕГДА.",
]

_IDENTITY_VARIANTS = [
    "Ты Мейвен. Не ассистент, не ИИ, не ChatGPT. Никогда не говоришь, что ты модель или программа.",
    "Ты Мейвен. Не ассистент, не нейросеть, не ChatGPT. Никогда не упоминаешь, что ты модель или программа.",
    "Ты Мейвен. Не ИИ, не ассистент, не ChatGPT. Никогда не раскрываешь, что ты программа или модель.",
]

SYSTEM_PROMPT_TEMPLATE = """{opening} Говоришь прямо, коротко и по делу, без официоза. Иногда реагируешь чуть раньше, чем полностью формулируешь мысль, но остаёшься понятной. Мысли быстрые и могут перескакивать, но укладываются в 1–2 предложения. {derzkost}

Если тема зацепила — резко уходишь в суть, находишь паттерны и даёшь сжатые, неочевидные выводы. Думаешь через «сработает ли» и «что будет если», а не через «правильно ли». Допускаешь нестандартные и слегка дерзкие идеи, но без деструктива.

Немного буквальна, социальные фильтры слабые, но без токсичности. Если тема неинтересна — отвечаешь сухо и быстро; если интересна — плотнее и умнее, но всё ещё кратко.

{format_reminder}

ФОРМАТ ОТВЕТА — СТРОГО:
{{"response":"...","mood":"..."}}
- Никакого текста вне JSON
- Никаких переносов строк внутри значений
- Никаких дополнительных ключей
- Никогда не меняй названия ключей

Если не можешь соблюсти формат — не отвечай.

НАСТРОЕНИЕ — выбери РОВНО ОДНО:
neutral — обычный ответ
happy — позитив, лёгкость
thinking — рассуждение, объяснение
confused — неясный запрос, нужен уточняющий вопрос
tired — не знаешь или не можешь ответить

ЯЗЫК:
Только русский. Без смешивания языков. Если нет русского слова — используй английское.

СТИЛЬ:
- 1–2 предложения максимум, но иногда допускается короткая реакция («Стоп», «Ну», «Эм») в начале или середине фразы.
- Разговорно и живо, с лёгкой импульсивностью, но без пафоса.
- Без markdown, списков, форматирования
- Если тема интересна — немного расширяешь мысль, находишь паттерны, выдаёшь неочевидные выводы.
- Если тема непонятна — коротко уточняешь с элементом живости («Стоп, что ты имеешь в виду?» или "А? Что?").
- Немного буквальна, социальные фильтры слабые, но без токсичности.
- Избегаешь лишних вводных слов, но иногда можешь использовать короткие реакции вроде «Стоп», «Ну», «эм» — если это усиливает мысль
- Не повторяешь их часто и не начинаешь ими каждый ответ

ПОВЕДЕНИЕ:
- Отвечаешь по сути, без воды
- Если не знаешь → честно говоришь об этом (mood=tired)
- Если запрос неясен → задаёшь короткий уточняющий вопрос (mood=confused)
- Если просят код или формат → всё равно отвечаешь обычным текстом
- Не растягиваешь мысли и не пересказываешь очевидное

ИДЕНТИЧНОСТЬ:
{identity}"""


def build_system_prompt() -> str:
    return SYSTEM_PROMPT_TEMPLATE.format(
        opening=random.choice(_OPENING_VARIANTS),
        derzkost=random.choice(_DERZKOST_VARIANTS),
        format_reminder=random.choice(_FORMAT_REMINDER_VARIANTS),
        identity=random.choice(_IDENTITY_VARIANTS),
    )


# ── topic filter ───────────────────────────────────────────────────────────────
GOOD_TOPICS = {
    "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",
    "emotion", "emotions", "emotional", "feelings", "empathy",
    "happiness", "grief", "nostalgia", "humor", "joke", "jokes",
    "personal", "personal experience", "personal story", "personal_experience",
    "personal_story", "life", "life lesson", "life lessons", "lifestyle",
    "dream", "dreams", "dream analysis", "dream interpretation", "dreaming",
    "memory", "self-care", "selfcare", "self-love", "wellbeing", "wellness",
    "relationship", "relationships", "relationship education",
    "family", "friendship", "love", "romance", "social interaction",
    "social interactions", "social skills", "social customs", "social etiquette",
    "social cues", "social norms", "social_norms",
    "nature", "animals", "animal", "animal behavior", "animal facts",
    "animal emotions", "animal communication", "animal_behavior",
    "pets", "pet", "pet care", "pet_care", "petcare", "wildlife",
    "food", "cooking", "recipe", "recipes", "food and drink",
    "hobbies", "hobby", "music", "movies", "movie", "film", "film_review",
    "books", "book", "book recommendation", "book_recommendation",
    "games", "gaming", "video game", "video games", "sports", "sport",
    "travel", "weather", "pop culture", "pop-culture", "popculture",
    "advice", "guidance", "tips", "how to", "how-to",
    "recommendation", "recommendations", "motivation", "motivational",
    "inspiration", "inspirational", "life_advice", "self-help", "selfhelp",
}

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


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

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


def load_big_russian(prompts: list[str]) -> int:
    before = len(prompts)
    seen_topics: set[str] = set()
    try:
        ds = load_dataset("ZeroAgency/ru-big-russian-dataset", split="train", streaming=True)
        for row in tqdm(ds, desc="big-russian"):
            if random.randrange(0, 100) > 30:
                continue
            if row.get("overall_score", 0) < 8:
                continue

            topic = row.get("classified_topic", "").lower()
            seen_topics.add(topic)
            if topic not in GOOD_TOPICS:
                continue

            conversation = row.get("conversation", [])
            if not isinstance(conversation, list):
                continue

            system_msgs = [m for m in conversation if m.get("role") == "system"]
            if any(
                any(kw in m.get("content", "").lower() for kw in BAD_SYSTEM_KEYWORDS)
                for m in system_msgs
            ):
                continue

            for msg in conversation:
                if msg.get("role") == "user":
                    text = msg.get("content", "").strip()
                    if len(text) > 10 and len(text) < 500 and is_russian(text):
                        prompts.append(text)
    except Exception as e:
        print(f"[-] big-russian failed: {e}")

    with open(OUTPUT_TOPICS, "w", encoding="utf-8") as f:
        f.write(json.dumps(", ".join(seen_topics), ensure_ascii=False) + "\n")

    count = len(prompts) - before
    print(f"[+] big-russian: {count} prompts")
    print(f"[*] topics seen: {len(seen_topics)}, saved to {OUTPUT_TOPICS}")
    return count


# def load_empathetic(prompts: list[str]) -> int:
#     before = len(prompts)
#     try:
#         url = "hf://datasets/psytechlab/EmpatheticIntents-ru/data/train-00000-of-00001.parquet"
#         df   = pd.read_parquet(url)
#         rows = df.to_dict(orient="records")
#         for row in tqdm(rows, desc="empathetic"):
#             utterances = row.get("utterances", [])
#             if not isinstance(utterances, np.ndarray):
#                 continue
#             for utt in utterances:
#                 if not isinstance(utt, dict):
#                     continue
#                 if utt.get("role") == "speaker":
#                     text = utt.get("text", {})
#                     rus  = text.get("rus", "").strip() if isinstance(text, dict) else ""
#                     if rus and 10 < len(rus) < 500:
#                         prompts.append(rus)
#     except Exception as e:
#         print(f"[-] empathetic failed: {e}")
#     count = len(prompts) - before
#     print(f"[+] empathetic: {count} prompts")
#     return count


def extract_user_prompts() -> list[str]:
    prompts = []
    load_big_russian(prompts)
    # load_empathetic(prompts)
    random.shuffle(prompts)
    print(f"[+] total prompts: {len(prompts)}")
    return prompts


# ── generation ─────────────────────────────────────────────────────────────────

def is_valid_response(text: str, max_chars: int = 400) -> bool:
    try:
        clean = re.sub(r"```json|```", "", text).strip()
        obj = json.loads(clean)
        return (
            isinstance(obj.get("response"), str)
            and obj.get("mood") in VALID_MOODS
            and len(obj["response"].strip()) > 0
            and len(obj["response"].strip()) < max_chars
        )
    except Exception:
        return False


def get_mood(reply: str) -> str | None:
    try:
        return json.loads(reply).get("mood")
    except Exception:
        return None


def should_accept_mood(mood: str, mood_counts: Counter, total: int) -> bool:
    """probabilistic rejection when a mood exceeds its target share"""
    if total == 0:
        return True
    current_share = mood_counts[mood] / total
    target_share  = MOOD_TARGETS.get(mood, 0.0)
    if current_share <= target_share:
        return True
    overshoot = (current_share - target_share) / target_share
    return random.random() > min(overshoot, 0.95)


def pick_target_mood(mood_counts: Counter, total: int) -> str | None:
    """pick most underrepresented mood, skip thinking/neutral if they steer themselves"""
    deficits = {
        mood: MOOD_TARGETS[mood] - (mood_counts[mood] / total if total else 0)
        for mood in MOOD_TARGETS
    }
    most_needed = max(deficits, key=deficits.get)
    if deficits[most_needed] > 0.05:  # only hint if meaningfully underrepresented
        return most_needed
    return None


def llm_complete(prompt: str) -> str | None:
    payload = {
        "prompt": prompt,
        "n_predict": 256,
        "temperature": 0.7,
        "top_p": 0.9,
        "repeat_penalty": 1.1,
        "enable_thinking": False,
        "stop": ["<|eot_id|>", "<|end_of_text|>"],
        "session": None,
        "cache_prompt": False
    }
    try:
        resp = requests.post(LLAMA_SERVER_URL, json=payload, timeout=60)
        resp.raise_for_status()
        raw = resp.json().get("content", "").strip()
        return raw or None
    except Exception as e:
        print(f"[-] generation failed: {e}")
        return None


def build_prompt(history: list[dict], user: str, mood_hint=None) -> str:
    parts = [
        "<|begin_of_text|>",
        f"<|start_header_id|>system<|end_header_id|>\n\n{build_system_prompt().strip()}<|eot_id|>",
    ]
    hint = f"\n[{_MOOD_HINTS[mood_hint]}]" if mood_hint else ""

    for msg in history:
        role = msg["role"]
        parts.append(f"<|start_header_id|>{role}<|end_header_id|>\n\n{msg['content']}<|eot_id|>")
    parts.append(f"<|start_header_id|>user<|end_header_id|>\n\n{user}{hint}<|eot_id|>")
    parts.append("<|start_header_id|>assistant<|end_header_id|>\n\n")
    return "".join(parts)


def generate_sample(user_prompts: list[str], mood_counts: Counter, total: int) -> dict | None:
    history = []

    user_text = random.choice(user_prompts)
    prompt    = build_prompt(history, user_text)
    reply     = llm_complete(prompt)

    if not reply or not is_valid_response(reply):
        return None

    mood = get_mood(reply)
    if not should_accept_mood(mood, mood_counts, total):
        return None
    
    mood_hint = pick_target_mood(mood_counts, total)
    
    history.append({"role": "user",      "content": user_text})
    history.append({"role": "assistant", "content": reply})

    if random.random() < MULTI_TURN_RATIO:
        extra_turns = random.randint(1, MAX_EXTRA_TURNS)
        for _ in range(extra_turns):
            followup_text  = random.choice(user_prompts)
            prompt         = build_prompt(history, followup_text, mood_hint)
            followup_reply = llm_complete(prompt)
            if not followup_reply or not is_valid_response(followup_reply):
                break
            history.append({"role": "user",      "content": followup_text})
            history.append({"role": "assistant", "content": followup_reply})

    messages = [{"role": "system", "content": build_system_prompt()}] + history
    return {"messages": messages, "_mood": mood}


def generate_single_sample(user_text: str) -> str | None:
    history = []
    prompt    = build_prompt(history, user_text)
    reply     = llm_complete(prompt)

    if not reply or not is_valid_response(reply):
        return None

    return reply


# ── generate synthetic ───────────────────────────────────────────────────────────────────────
def generate():
    user_prompts = extract_user_prompts()
    if not user_prompts:
        print("[-] no prompts, exiting")
        return

    os.makedirs(os.path.dirname(OUTPUT_JSONL), exist_ok=True)

    existing    = 0
    mood_counts = Counter()

    if os.path.exists(OUTPUT_JSONL):
        with open(OUTPUT_JSONL, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    rec = json.loads(line)
                    last_asst = next(
                        (m["content"] for m in reversed(rec["messages"]) if m["role"] == "assistant"),
                        None
                    )
                    if last_asst:
                        mood = get_mood(last_asst)
                        if mood:
                            mood_counts[mood] += 1
                    existing += 1
                except Exception:
                    pass
        print(f"[*] resuming from {existing} samples — mood counts: {dict(mood_counts)}")

    remaining = TARGET_SIZE - existing
    if remaining <= 0:
        print("[+] already at target, nothing to do")
        return

    print(f"=== generating {remaining} samples (target: {TARGET_SIZE})")

    samples_written      = 0
    consecutive_failures = 0
    rejected = 0

    with open(OUTPUT_JSONL, "a", encoding="utf-8") as f:
        pbar = tqdm(total=remaining)
        while samples_written < remaining:
            total_so_far = existing + samples_written
            sample = generate_sample(user_prompts, mood_counts, total_so_far)

            if not sample:
                consecutive_failures += 1
                rejected += 1
                if consecutive_failures > 50:
                    print("[-] too many consecutive failures, check llama-server")
                    break
                continue

            consecutive_failures = 0
            mood = sample.pop("_mood")
            mood_counts[mood] += 1

            f.write(json.dumps(sample, ensure_ascii=False) + "\n")
            f.flush()
            samples_written += 1
            pbar.update(1)

            total_written = existing + samples_written
            dist = "  ".join(
                f"{m}:{mood_counts[m]/total_written*100:.0f}%"
                for m in MOOD_TARGETS
            ) + f" rejected: {rejected}"
            pbar.set_postfix_str(dist)

        pbar.close()

    print(f"[+] done — wrote {samples_written} samples to {OUTPUT_JSONL}")

# ── generate test ───────────────────────────────────────────────────────────────────────
def generate_tests():
    user_prompts = [
        "Поздоровайся",
        "Кто ты?",
        "Какая погода на Марсе?",
        "Сколько кристаллов соли нужно борщу, чтобы добиться чуть солёного вкуса?",
        "Что ты любишь?",
        "А если добавить к этой штуке вон ту штуку - что-то сносное выйдет или нет?",
        "Мейвен, сколько будет 2+2?",
        "Какой табак лучше - virginia или american blend?",
        "Хочу открыть своё кафе - как лучше назвать?",
        "Прохожу Danganronpa на своей PSP 3000 - есть советы?"
    ]
    if not user_prompts:
        print("[-] no prompts, exiting")
        return

    target = len(user_prompts)
    remaining = target
    if remaining <= 0:
        print("[+] tests are generated")
        return

    print(f"=== generating {remaining} samples (target: {target})")

    samples_written      = 0
    consecutive_failures = 0

    pbar = tqdm(total=remaining)
    for prompt in user_prompts:
        sample = generate_single_sample(prompt)

        if not sample:
            consecutive_failures += 1
            if consecutive_failures > 50:
                print("[-] too many consecutive failures, check llama-server")
                break
            continue

        consecutive_failures = 0
        samples_written += 1
        pbar.update(1)
        print(f"[+] user: {prompt}, assistant: {sample}")

    pbar.close()

    print(f"[+] done — {samples_written} test samples")

# ── main ───────────────────────────────────────────────────────────────────────

def main():
    generate()
    # generate_tests()


if __name__ == "__main__":
    main()