#!/usr/bin/env python3
"""
Generate synthetic dataset for Generic vs Semantic classifier using Ollama (llama3.1:8b).

Generates 4 categories:
  - en_generic:  English generic queries
  - en_semantic: English semantic queries
  - hi_generic:  Hindi generic queries (Devanagari)
  - hi_semantic: Hindi semantic queries (Devanagari)

Each category targets TOTAL_PER_CATEGORY examples (default 3000).
Generation is resumable — it appends to existing JSONL files.

Usage:
    python3 scripts/generate_dataset.py [--category en_generic]
    python3 scripts/generate_dataset.py              # all categories
"""

import json
import os
import re
import sys
import time
import argparse

import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import (
    CATEGORIES, TOTAL_PER_CATEGORY, BATCH_SIZE_GEN, MAX_CONCURRENT,
    OLLAMA_URL, OLLAMA_MODEL, RAW_DIR
)

os.makedirs(RAW_DIR, exist_ok=True)


def count_existing(filepath: str) -> int:
    """Count lines (examples) in an existing JSONL file."""
    if not os.path.exists(filepath):
        return 0
    with open(filepath) as f:
        return sum(1 for _ in f)


def build_prompt(category_key: str) -> str:
    """Build a prompt for the given category that asks for exactly BATCH_SIZE_GEN examples."""
    info = CATEGORIES[category_key]
    lang = info["lang"]
    label = info["label"]

    # Language-specific instructions
    if lang == "Hindi":
        lang_instructions = """- Write ALL queries in Devanagari script (Hindi), NOT transliterated Hindi.
- Use conversational Hindi, not formal/literary Hindi.
- Include natural particles like "ही", "भी", "तो", "ना", "जी".
- Use common Hindi interjections: "अच्छा", "हाँ", "नहीं", "है ना", "अरे". """
    else:
        lang_instructions = """- Write in natural, conversational English.
- Cover different registers: casual, polite, formal, technical."""

    # Category-specific definitions and examples
    if label == "GENERIC":
        gen_examples = (
            f"Example GENERIC {lang} queries:\n"
            f'        - hello / namaste\n'
            f'        - stop talking / chup raho\n'
            f'        - what time is it / kya samay hua hai\n'
            f'        - thanks / shukriya\n'
            f'        - okay got it / theek hai samajh gaya\n'
            f'        - tell me a joke / koi chutkula sunao\n'
            f'        - i see / achha\n'
            f'        - yes please continue / haan ji kripya jari rakhein\n'
            f'        - how are you / aap kaise hain\n'
            f'        - never mind / koi baat nahi\n'
        )
        definition = (
            "GENERIC queries have NO durable knowledge value. They are:\n"
            "- Social rituals: greetings, thanks, apologies, pleasantries\n"
            "- Commands/Controls: start, stop, pause, go back, repeat\n"
            "- Simple affirmations/negations: yes, no, okay, hmm, got it\n"
            '- Simple time/date/weather queries ("what time is it")\n'
            "- Fillers and backchanneling: well, so, anyway, i see, right\n"
            '- Transactional: "please repeat", "speak slower", "tell me a joke"\n'
            '- Interaction management: "im done", "thats all", "go ahead"\n'
            '- Unanswerable/meta: "i dont know", "what do you mean", "can you hear me"\n'
        )
    else:  # SEMANTIC
        gen_examples = (
            f"Example SEMANTIC {lang} queries:\n"
            f"  SHORT (3-7 words) standalone semantic statements:\n"
            f'        - my name is John / mera naam Ravi hai\n'
            f'        - I am a doctor / main doctor hoon\n'
            f'        - I love spicy food / mujhe masaledar khana pasand hai\n'
            f'        - my sister is a teacher / meri behen teacher hai\n'
            f'        - I live in Delhi / main Dilli mein rehta hoon\n'
            f'        - I work at Google / main Google mein kaam karta hoon\n'
            f'        - my favorite color is blue / mera pasandida rang nila hai\n'
            f'        - I have two cats / mere paas do billiyan hain\n'
            f'        - I am learning guitar / main guitar seekh raha hoon\n'
            f'  LONGER (8-20 words) compound semantic statements:\n'
            f'        - my name is John and I live in Mumbai / mera naam Ravi hai aur main Mumbai mein rehta hoon\n'
            f'        - I love spicy food but I am allergic to peanuts / mujhe masaledar khana pasand hai lekin mujhe moongphali se allergy hai\n'
            f'        - my sister is a doctor in Delhi / meri behen Dilli mein doctor hai\n'
            f'        - I am planning to start learning guitar next month / main agle mahine guitar seekhna shuru karne wala hoon\n'
            f'        - remember I said I am allergic to peanuts / yaad hai maine kaha tha mujhe moongphali se allergy hai\n'
            f'        - my favorite restaurant is the Italian place on Church Street / mera pasandida restaurant Church Street par Italian jagah hai\n'
        )
        definition = (
            "SEMANTIC queries contain durable, storable information. They are:\n"
            "- Personal facts: name, age, location, profession, education, background\n"
            "- Preferences and tastes: likes, dislikes, favorites, habits\n"
            "- Relationships: family, friends, colleagues, their attributes\n"
            "- Detailed descriptions of events, people, places, objects\n"
            "- Complex questions that require retrieval of past context\n"
            '- Explicit memory references: "remember I told you about...", "as I said before..."\n'
            '- Plans, intentions, goals: "Im planning to visit Japan next spring"\n'
            '- OPINIONS WITH REASONING: "I think dark chocolate is better because..."\n'
            '- Knowledge queries that reveal user context: "How long does it take to get to Bangalore?"\n'
            "  (These reveal the user's location/context even though they are phrased as questions)\n"
        )

    lang_code = "hi" if lang == "Hindi" else "en"
    prompt = (
        f"You are generating a synthetic training dataset for a binary classifier. "
        f"The classifier categorizes user queries as GENERIC (no durable knowledge) "
        f"or SEMANTIC (contains storable facts, preferences, relationships, context).\n\n"
        f"TASK: Generate {BATCH_SIZE_GEN} realistic {lang} user queries. "
        f"EVERY query must be labeled \"{label}\".\n\n"
        f"{lang_instructions}\n\n"
        f"{definition}\n\n"
        f"{gen_examples}\n\n"
        f"CRITICAL RULES:\n"
        f'1. Every query MUST have label = "{label}" - no mix of labels.\n'
        f"2. Output ONLY valid JSONL - one JSON object per line, nothing else.\n"
        f'3. Each line format: {{\"text\": \"<the query>\", "language\": \"{lang_code}\", "label\": "{label}"}}\n'
        f"4. Queries must be diverse: vary the patterns, structures, and lengths (2 to 20 words).\n"
    )

    if label == "SEMANTIC":
        prompt += (
            f"5. IMPORTANT - 40% of your examples MUST be SHORT (3-7 words) standalone statements "
            f"containing exactly one fact/preference. The remaining 60% can be longer compound sentences.\n"
        )
    else:
        prompt += (
            f"5. Make them sound like real voice assistant queries, not textbook sentences.\n"
        )

    prompt += (
        f"6. NO markdown, NO code fences, NO explanation, NO numbering.\n\n"
        f"Now generate {BATCH_SIZE_GEN} examples, one per line:"
    )

    return prompt


def parse_jsonl_from_response(content: str) -> list[dict]:
    """Parse JSONL from the model response, handling common formatting issues."""
    examples = []
    for line in content.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        # Remove markdown code fences
        if line.startswith("```"):
            continue
        if line == '```':
            continue

        # Try direct JSON parse
        try:
            obj = json.loads(line)
            if "text" in obj and "label" in obj:
                obj["label"] = obj["label"].strip().upper()
                examples.append(obj)
                continue
        except json.JSONDecodeError:
            pass

        # Try to find JSON within the line
        match = re.search(r'\{[^}]*"text"[^}]*"label"[^}]*\}', line)
        if match:
            try:
                obj = json.loads(match.group())
                if "text" in obj and "label" in obj:
                    obj["label"] = obj["label"].strip().upper()
                    examples.append(obj)
            except json.JSONDecodeError:
                pass

    return examples


def generate_batch(category: str) -> list[dict]:
    """Generate one batch of examples from Ollama."""
    prompt = build_prompt(category)

    payload = {
        "model": OLLAMA_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
        "options": {
            "temperature": 0.85,
            "top_p": 0.95,
            "num_predict": 4096,
        }
    }

    try:
        resp = requests.post(OLLAMA_URL, json=payload, timeout=300)
        resp.raise_for_status()
        content = resp.json()["message"]["content"]
        examples = parse_jsonl_from_response(content)
        return examples
    except requests.exceptions.Timeout:
        print(f"    [TIMEOUT] Batch generation timed out")
        return []
    except Exception as e:
        print(f"    [ERROR] {e}")
        return []


def generate_category(category: str):
    """Generate TOTAL_PER_CATEGORY examples for one category using concurrent batches."""
    filepath = os.path.join(RAW_DIR, f"{category}.jsonl")
    existing = count_existing(filepath)
    needed = TOTAL_PER_CATEGORY - existing

    if needed <= 0:
        print(f"  [SKIP] {category}: already has {existing} examples (target {TOTAL_PER_CATEGORY})")
        return

    print(f"  [GEN]  {category}: {existing} existing, {needed} more needed")

    generated_count = existing
    pbar = tqdm(total=TOTAL_PER_CATEGORY, initial=existing, desc=f"{category:15s}", unit="ex", smoothing=0.1)

    # Calculate how many batches we need (with a safety margin)
    batches_to_submit = needed // BATCH_SIZE_GEN + 3  # overshoot slightly
    submitted = 0

    with ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor:
        # Submit initial batches
        futures = {}
        initial_count = min(MAX_CONCURRENT, batches_to_submit)
        for _ in range(initial_count):
            future = executor.submit(generate_batch, category)
            futures[future] = True
            submitted += 1

        # Process as they complete, submitting more to maintain throughput
        while futures and generated_count < TOTAL_PER_CATEGORY:
            for future in as_completed(futures, timeout=120):
                break  # just get one

            try:
                examples = future.result()
                if examples:
                    with open(filepath, "a") as fh:
                        for ex in examples:
                            fh.write(json.dumps(ex, ensure_ascii=False) + "\n")
                    generated_count += len(examples)
                    pbar.update(len(examples))
            except Exception as e:
                print(f"    [ERROR] Batch failed: {e}")

            del futures[future]

            # Submit replacement if we haven't submitted all needed
            if submitted < batches_to_submit and generated_count < TOTAL_PER_CATEGORY * 1.1:
                new_future = executor.submit(generate_batch, category)
                futures[new_future] = True
                submitted += 1

    pbar.close()
    final_count = count_existing(filepath)
    print(f"  [DONE] {category}: {final_count} examples")


def main():
    parser = argparse.ArgumentParser(description="Generate Generic vs Semantic dataset")
    parser.add_argument("--category", "-c", choices=list(CATEGORIES.keys()) + ["all"], default="all",
                       help="Category to generate (default: all)")
    args = parser.parse_args()

    categories = list(CATEGORIES.keys()) if args.category == "all" else [args.category]

    print(f"=" * 60)
    print(f"Generic vs Semantic Dataset Generator")
    print(f"Target: {TOTAL_PER_CATEGORY} per category × {len(categories)} = {TOTAL_PER_CATEGORY * len(categories)} total")
    print(f"Ollama model: {OLLAMA_MODEL}")
    print(f"Concurrent: {MAX_CONCURRENT} workers, {BATCH_SIZE_GEN} per batch")
    print(f"Output: {RAW_DIR}/")
    print(f"=" * 60)

    for category in categories:
        generate_category(category)

    # Summary
    print(f"\n{'=' * 60}")
    print(f"Generation Complete — Summary:")
    print(f"{'=' * 60}")
    total = 0
    for category in categories:
        filepath = os.path.join(RAW_DIR, f"{category}.jsonl")
        count = count_existing(filepath)
        lang = CATEGORIES[category]["lang"]
        label = CATEGORIES[category]["label"]
        print(f"  {lang:8s} {label:8s}: {count:5d}")
        total += count
    print(f"  {'TOTAL':18s}: {total}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()