"""
scripts/04_generate_synth_dialogues.py
Generate synthetic (pictogram_sequence → target_utterance) training pairs
for each persona using Gemini Pro or Claude API.

Usage:
    python scripts/04_generate_synth_dialogues.py --api gemini --persona arjun
    python scripts/04_generate_synth_dialogues.py --api claude --all
"""

import argparse
import json
import logging
import os
import random
import time
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

SYNTH_DIR = Path("data/synth")
TARGET_PER_PERSONA = 3000
BATCH_SIZE = 50  # utterances per API call

# Meal and time contexts to diversify examples
CONTEXTS = [
    "breakfast", "lunch", "dinner", "school", "bedtime", "morning",
    "physiotherapy", "playing", "watching TV", "visiting temple/church/mosque",
    "shopping with family", "bath time", "doctor visit", "talking to grandparent",
]

# AAC pictogram categories for sequences
PICTOGRAM_CATEGORIES = {
    "people": ["mother", "father", "sister", "brother", "teacher", "doctor", "friend"],
    "food": ["water", "milk", "rice", "bread", "fruit", "medicine", "biscuit"],
    "actions": ["want", "give", "go", "come", "stop", "help", "eat", "drink", "sleep", "play"],
    "feelings": ["happy", "sad", "pain", "tired", "scared", "angry", "love"],
    "objects": ["tablet", "phone", "TV", "book", "toy", "cup", "toilet", "bed", "school"],
    "places": ["home", "school", "hospital", "outside", "temple", "market"],
    "time": ["now", "later", "today", "morning", "night"],
    "modifiers": ["more", "no", "yes", "big", "hot", "cold", "mine"],
}


def sample_pictogram_sequence(length: int = None) -> list[str]:
    """Generate a realistic pictogram tap sequence of 1-4 symbols."""
    if length is None:
        length = random.choices([1, 2, 3, 4], weights=[0.1, 0.35, 0.4, 0.15])[0]
    sequence = []
    cats = list(PICTOGRAM_CATEGORIES.keys())
    for _ in range(length):
        cat = random.choice(cats)
        symbol = random.choice(PICTOGRAM_CATEGORIES[cat])
        sequence.append(symbol)
    return list(dict.fromkeys(sequence))  # deduplicate while preserving order


def build_gemini_prompt(persona_dict: dict, context: str, sequences: list[list[str]]) -> str:
    """Build the prompt for Gemini to generate utterances."""
    lang_str = f"{persona_dict['primary_language']} (primary), {', '.join(persona_dict['secondary_languages'])}"
    sequences_str = "\n".join(
        f"{i+1}. [{', '.join(seq)}]" for i, seq in enumerate(sequences)
    )
    return f"""You are generating training data for an AAC (Augmentative & Alternative Communication) app for children with cerebral palsy in India.

Child profile:
- Name: {persona_dict['name']}, Age: {persona_dict['age']}
- City: {persona_dict['city']}, {persona_dict['state']}
- CP type: {persona_dict['cp_type']} ({persona_dict['severity']})
- Languages: {lang_str}
- Code-switching: {persona_dict['code_switch_pattern']}
- Favourite foods: {', '.join(persona_dict['favourite_foods'])}
- Favourite activities: {', '.join(persona_dict['favourite_activities'])}

Context: {context}

For each pictogram sequence below, generate the MOST LIKELY natural sentence this child wants to communicate. Rules:
1. Use the child's actual language mix (NOT just English)
2. Keep sentences short (5-15 words), age-appropriate
3. Be specific to the child's life (use their favourite foods/activities when relevant)
4. Also generate 2 ALTERNATIVE sentences (plausible but less likely)
5. Also generate 1 NEGATIVE example (a sentence that would be WRONG to output — e.g. inappropriate for a child, or completely off-context)
6. Output ONLY valid JSON, no preamble

Pictogram sequences:
{sequences_str}

Output format (one JSON object per line, {len(sequences)} lines total):
{{"seq": ["symbol1", "symbol2"], "target": "primary sentence", "alternatives": ["alt1", "alt2"], "negative": "wrong sentence", "context": "{context}"}}"""


def call_gemini(prompt: str, model: str = "gemini-1.5-pro") -> str:
    import google.generativeai as genai
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        raise EnvironmentError("GOOGLE_API_KEY not set")
    genai.configure(api_key=api_key)
    m = genai.GenerativeModel(model)
    response = m.generate_content(prompt)
    return response.text


def call_claude(prompt: str, model: str = "claude-opus-4-6") -> str:
    import anthropic
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise EnvironmentError("ANTHROPIC_API_KEY not set")
    client = anthropic.Anthropic(api_key=api_key)
    msg = client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{"role": "user", "content": prompt}]
    )
    return msg.content[0].text


def parse_api_response(text: str) -> list[dict]:
    """Parse newline-delimited JSON from the API response."""
    records = []
    # Strip markdown fences if present
    text = text.strip()
    if text.startswith("```"):
        text = text.split("\n", 1)[-1].rsplit("```", 1)[0]

    for line in text.strip().split("\n"):
        line = line.strip()
        if not line or not line.startswith("{"):
            continue
        try:
            obj = json.loads(line)
            if "target" in obj and "seq" in obj:
                records.append(obj)
        except json.JSONDecodeError:
            continue
    return records


def generate_for_persona(
    persona_id: str,
    api: str = "gemini",
    target: int = TARGET_PER_PERSONA,
) -> list[dict]:
    import sys
    sys.path.insert(0, "src")
    from ankahi.data.persona import PERSONAS

    persona = PERSONAS[persona_id]
    persona_dict = persona.to_dict()

    out_path = SYNTH_DIR / f"persona_{persona_id}.jsonl"
    SYNTH_DIR.mkdir(parents=True, exist_ok=True)

    # Resume if partially done
    existing = []
    if out_path.exists():
        with open(out_path) as f:
            existing = [json.loads(l) for l in f if l.strip()]
        log.info(f"Resuming: {len(existing)} records already in {out_path}")

    all_records = existing[:]
    needed = target - len(all_records)

    if needed <= 0:
        log.info(f"Persona {persona_id} already has {len(all_records)} records. Skipping.")
        return all_records

    log.info(f"Generating {needed} records for persona: {persona.name}")

    call_fn = call_gemini if api == "gemini" else call_claude
    contexts = CONTEXTS * ((needed // (len(CONTEXTS) * BATCH_SIZE)) + 2)

    batch_num = 0
    with open(out_path, "a", encoding="utf-8") as f:
        while len(all_records) - len(existing) < needed:
            context = contexts[batch_num % len(contexts)]
            batch_seqs = [sample_pictogram_sequence() for _ in range(BATCH_SIZE)]
            prompt = build_gemini_prompt(persona_dict, context, batch_seqs)

            try:
                raw = call_fn(prompt)
                records = parse_api_response(raw)
                for r in records:
                    r["persona_id"] = persona_id
                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
                all_records.extend(records)
                log.info(f"  Batch {batch_num}: got {len(records)} records (total {len(all_records)})")
            except Exception as e:
                log.error(f"  API call failed (batch {batch_num}): {e}")
                time.sleep(5)

            batch_num += 1
            time.sleep(1.5)  # rate limiting

    log.info(f"Done: {len(all_records)} records for {persona.name}")
    return all_records


def generate_shared_base(api: str = "gemini", n: int = 1500) -> list[dict]:
    """Generate a shared base corpus not tied to any persona."""
    out_path = SYNTH_DIR / "base_shared.jsonl"
    if out_path.exists():
        with open(out_path) as f:
            existing = [json.loads(l) for l in f if l.strip()]
        if len(existing) >= n:
            log.info(f"Base corpus already done ({len(existing)} records).")
            return existing

    # Use a generic Indian child persona
    generic_persona = {
        "name": "child", "age": 8, "city": "India", "state": "India",
        "cp_type": "spastic", "severity": "moderate",
        "primary_language": "hi", "secondary_languages": ["en"],
        "code_switch_pattern": "Hindi base with English nouns",
        "favourite_foods": ["rice", "dal", "roti"], "favourite_activities": ["playing", "TV"],
    }
    call_fn = call_gemini if api == "gemini" else call_claude
    records = []
    with open(out_path, "w", encoding="utf-8") as f:
        while len(records) < n:
            context = random.choice(CONTEXTS)
            seqs = [sample_pictogram_sequence() for _ in range(BATCH_SIZE)]
            prompt = build_gemini_prompt(generic_persona, context, seqs)
            try:
                raw = call_fn(prompt)
                batch = parse_api_response(raw)
                for r in batch:
                    r["persona_id"] = "base"
                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
                records.extend(batch)
                log.info(f"  Base: {len(records)}/{n}")
            except Exception as e:
                log.error(f"  Base batch error: {e}")
                time.sleep(5)
            time.sleep(1.5)
    return records


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--api", choices=["gemini", "claude"], default="gemini")
    parser.add_argument("--persona", type=str, default=None,
                        help="Single persona ID (ananya/arjun/priya/rohan/zara)")
    parser.add_argument("--all", action="store_true", help="Generate for all personas + base")
    parser.add_argument("--base-only", action="store_true")
    parser.add_argument("--target", type=int, default=TARGET_PER_PERSONA)
    args = parser.parse_args()

    if args.base_only:
        generate_shared_base(args.api)
        return

    if args.all:
        generate_shared_base(args.api)
        for pid in ["ananya", "arjun", "priya", "rohan", "zara"]:
            generate_for_persona(pid, args.api, args.target)
    elif args.persona:
        generate_for_persona(args.persona, args.api, args.target)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()