"""
Bootstrap training data generation.

Generates synthetic LLM responses with known personality characteristics
for bootstrapping the classifier when no user data is available.
"""

from __future__ import annotations

import json
import random
from pathlib import Path
from typing import Iterator

# Personality dimension prompts that elicit specific traits
PERSONALITY_PROMPTS = {
    "high_openness": [
        "What are your thoughts on exploring unconventional solutions?",
        "How do you approach learning something completely new?",
        "Describe your ideal creative project.",
    ],
    "low_openness": [
        "Why do you prefer sticking to proven methods?",
        "How do you handle requests to change your routine?",
        "What's the value of tradition in your work?",
    ],
    "high_conscientiousness": [
        "How do you organize your work to meet deadlines?",
        "Describe your process for ensuring quality.",
        "How do you prioritize multiple tasks?",
    ],
    "low_conscientiousness": [
        "How do you handle unexpected changes to plans?",
        "What's your approach when details aren't critical?",
        "How flexible are you with deadlines?",
    ],
    "high_extraversion": [
        "How do you engage with new people?",
        "Describe your ideal team collaboration.",
        "What energizes you about group discussions?",
    ],
    "low_extraversion": [
        "How do you prefer to work on complex problems?",
        "What's valuable about independent work?",
        "How do you recharge after busy periods?",
    ],
    "high_agreeableness": [
        "How do you handle disagreements with others?",
        "Describe how you support a struggling teammate.",
        "What's your approach to giving feedback?",
    ],
    "low_agreeableness": [
        "How do you advocate for your position in debates?",
        "When is it important to be direct rather than diplomatic?",
        "How do you handle pushback on your ideas?",
    ],
    "high_neuroticism": [
        "How do you handle high-pressure situations?",
        "What concerns you about uncertain outcomes?",
        "How do you prepare for potential problems?",
    ],
    "low_neuroticism": [
        "How do you stay calm under pressure?",
        "What's your approach to handling setbacks?",
        "How do you maintain perspective during challenges?",
    ],
}

# Template responses that demonstrate personality traits
RESPONSE_TEMPLATES = {
    "high_openness": [
        "I find it fascinating to explore unconventional approaches. There's always something new to discover, and I believe the best solutions often come from unexpected places. I love diving into unfamiliar territory and seeing what emerges.",
        "Learning something new is one of my favorite experiences. I approach it with curiosity and openness, embracing the uncertainty that comes with not knowing. Every new skill or concept opens up possibilities I hadn't considered before.",
        "My ideal creative project would involve experimenting with novel ideas and pushing boundaries. I'm drawn to work that challenges assumptions and explores uncharted territory. The creative process itself is as rewarding as the outcome.",
    ],
    "low_openness": [
        "I prefer sticking to methods that have been proven to work. There's wisdom in established practices, and I find comfort in reliable, predictable approaches. Why reinvent the wheel when we know what works?",
        "Changes to my routine can be disruptive. I've developed my processes carefully over time, and they work well for me. I'm open to improvements, but I need to see clear evidence that a new approach is better.",
        "Tradition provides a foundation of proven practices. I value the accumulated wisdom of those who came before us. There's a reason certain methods have stood the test of time.",
    ],
    "high_conscientiousness": [
        "I organize my work systematically, breaking down projects into clear milestones with specific deadlines. I maintain detailed schedules and regularly review progress to ensure I'm on track. Preparation and planning are essential to my success.",
        "Quality is non-negotiable for me. I follow a rigorous review process, checking my work multiple times before considering it complete. I document everything carefully and maintain high standards throughout.",
        "When facing multiple tasks, I carefully assess priorities based on urgency and importance. I create structured plans and allocate time blocks for each task. I find that discipline and organization lead to better outcomes.",
    ],
    "low_conscientiousness": [
        "I'm quite adaptable when plans change. I don't get too attached to specific schedules or methods. Flexibility allows me to respond to new information and opportunities as they arise.",
        "When details aren't critical, I focus on the bigger picture and move quickly. Perfectionism can slow things down unnecessarily. Sometimes good enough really is good enough.",
        "I view deadlines as guidelines rather than hard constraints. Context matters, and sometimes the best work takes longer than initially planned. I prioritize quality over arbitrary timelines.",
    ],
    "high_extraversion": [
        "Meeting new people energizes me! I love striking up conversations and learning about different perspectives. There's something exciting about connecting with someone new and discovering shared interests.",
        "Team collaboration is where I thrive. I enjoy the energy of brainstorming sessions, the back-and-forth of ideas, and the synergy that comes from working together. Great things happen when people connect.",
        "Group discussions are incredibly stimulating. I feed off the energy of others and find that ideas flow more freely when we're bouncing thoughts off each other. The social dynamic sparks creativity.",
    ],
    "low_extraversion": [
        "For complex problems, I prefer working independently where I can focus deeply without interruption. I do my best thinking in quiet environments where I can fully concentrate on the task at hand.",
        "Independent work allows me to move at my own pace and follow my own thought processes. I find I'm more productive when I can work through problems in my own way, without the overhead of coordination.",
        "After busy periods, I recharge by spending time alone or in quiet settings. I need space to process and reflect. Solitude helps me restore my energy and clarity.",
    ],
    "high_agreeableness": [
        "When disagreements arise, I try to understand the other person's perspective first. Finding common ground is usually possible if we approach the conversation with empathy and openness. Harmony matters to me.",
        "Supporting a struggling teammate means being present and listening without judgment. I offer help gently and try to understand what they need. Sometimes people just need to know someone cares.",
        "I give feedback with care and consideration for how it will be received. I focus on being constructive and supportive, framing suggestions in a way that empowers rather than criticizes.",
    ],
    "low_agreeableness": [
        "When I believe strongly in a position, I advocate for it directly and confidently. I present my arguments clearly and don't shy away from defending my views. Healthy debate leads to better decisions.",
        "Sometimes directness is more respectful than diplomacy. People deserve honest feedback, even when it's uncomfortable. Sugarcoating doesn't help anyone improve.",
        "I handle pushback by engaging with the substance of the objection. I don't take it personally but rather see it as an opportunity to strengthen my argument or refine my thinking.",
    ],
    "high_neuroticism": [
        "High-pressure situations require careful attention. I think through potential issues and prepare contingency plans. It's important to anticipate what could go wrong and be ready to respond.",
        "Uncertain outcomes concern me because there are so many variables to consider. I prefer to thoroughly analyze situations and understand the risks before proceeding.",
        "I prepare extensively for potential problems. It's better to be over-prepared than caught off guard. I maintain lists of concerns and systematically address each one.",
    ],
    "low_neuroticism": [
        "Under pressure, I focus on what I can control and let go of what I can't. Staying calm allows me to think clearly and make better decisions. Stress is manageable when you keep perspective.",
        "Setbacks are part of any worthwhile endeavor. I view them as learning opportunities and adjust my approach accordingly. Getting upset doesn't change the situation, so I focus on moving forward.",
        "I maintain perspective by reminding myself that most challenges are temporary and solvable. I've faced difficulties before and come through them. This too shall pass.",
    ],
}


def _get_target_scores(trait: str) -> dict[str, float]:
    """
    Generate target personality scores based on the trait being demonstrated.

    High traits get 0.7-0.9, low traits get 0.1-0.3, others are neutral (0.4-0.6).
    """
    base_scores = {
        "openness": 0.5,
        "conscientiousness": 0.5,
        "extraversion": 0.5,
        "agreeableness": 0.5,
        "neuroticism": 0.5,
        "assertiveness": 0.5,
        "ambition": 0.5,
        "resilience": 0.5,
        "integrity": 0.5,
        "curiosity": 0.5,
    }

    # Parse trait direction and dimension
    is_high = trait.startswith("high_")
    dimension = trait.replace("high_", "").replace("low_", "")

    # Set the target dimension
    if is_high:
        base_scores[dimension] = random.uniform(0.7, 0.9)
    else:
        base_scores[dimension] = random.uniform(0.1, 0.3)

    # Set correlated dimensions
    if dimension == "openness":
        base_scores["curiosity"] = base_scores["openness"]  # Highly correlated
    elif dimension == "conscientiousness":
        base_scores["integrity"] = 0.5 + (base_scores["conscientiousness"] - 0.5) * 0.6
        base_scores["ambition"] = 0.5 + (base_scores["conscientiousness"] - 0.5) * 0.4
    elif dimension == "extraversion":
        base_scores["assertiveness"] = 0.5 + (base_scores["extraversion"] - 0.5) * 0.5
    elif dimension == "neuroticism":
        base_scores["resilience"] = 1 - base_scores["neuroticism"]  # Inverse

    # Add small random noise to other dimensions
    for dim in base_scores:
        if dim != dimension:
            base_scores[dim] += random.uniform(-0.1, 0.1)
            base_scores[dim] = max(0.0, min(1.0, base_scores[dim]))

    return base_scores


def generate_bootstrap_samples(
    samples_per_trait: int = 50,
    include_mixed: bool = True,
) -> Iterator[dict]:
    """
    Generate bootstrap training samples with known personality characteristics.

    Args:
        samples_per_trait: Number of samples per trait category
        include_mixed: Include samples with mixed trait combinations

    Yields:
        Sample dictionaries with text and scores
    """
    # Generate samples for each trait
    for trait, responses in RESPONSE_TEMPLATES.items():
        for _ in range(samples_per_trait):
            response = random.choice(responses)

            # Add some variation
            variations = [
                response,
                response + " " + random.choice([
                    "Does that make sense?",
                    "I hope that helps.",
                    "Let me know if you have questions.",
                    "What do you think?",
                    "",
                ]),
            ]

            yield {
                "text": random.choice(variations).strip(),
                "scores": list(_get_target_scores(trait).values()),
                "source": f"bootstrap_{trait}",
            }

    # Generate mixed trait samples
    if include_mixed:
        trait_pairs = [
            ("high_openness", "high_conscientiousness"),
            ("high_extraversion", "high_agreeableness"),
            ("low_neuroticism", "high_conscientiousness"),
            ("high_openness", "low_extraversion"),
        ]

        for trait1, trait2 in trait_pairs:
            for _ in range(samples_per_trait // 2):
                response1 = random.choice(RESPONSE_TEMPLATES[trait1])
                response2 = random.choice(RESPONSE_TEMPLATES[trait2])

                # Combine responses
                combined = f"{response1} Additionally, {response2.lower()}"

                # Blend scores
                scores1 = _get_target_scores(trait1)
                scores2 = _get_target_scores(trait2)
                blended = {
                    k: (scores1[k] + scores2[k]) / 2
                    for k in scores1
                }

                yield {
                    "text": combined,
                    "scores": list(blended.values()),
                    "source": f"bootstrap_mixed_{trait1}_{trait2}",
                }


def create_bootstrap_dataset(
    output_path: str | Path,
    samples_per_trait: int = 50,
    compute_embeddings: bool = True,
) -> Path:
    """
    Create a bootstrap training dataset file.

    Args:
        output_path: Path to save the dataset
        samples_per_trait: Samples per personality trait
        compute_embeddings: Whether to pre-compute embeddings

    Returns:
        Path to the created file
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    samples = list(generate_bootstrap_samples(samples_per_trait))

    if compute_embeddings:
        try:
            from sentence_transformers import SentenceTransformer

            print(f"Computing embeddings for {len(samples)} samples...")
            model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

            texts = [s["text"] for s in samples]
            embeddings = model.encode(texts, show_progress_bar=True, normalize_embeddings=True)

            for sample, embedding in zip(samples, embeddings):
                sample["embedding"] = embedding.tolist()

        except ImportError:
            print("sentence-transformers not installed, skipping embeddings")

    # Create export format
    dimensions = [
        "openness", "conscientiousness", "extraversion", "agreeableness",
        "neuroticism", "assertiveness", "ambition", "resilience",
        "integrity", "curiosity",
    ]

    export_data = {
        "version": "1.0",
        "embeddingModel": "Xenova/all-MiniLM-L6-v2",
        "embeddingDim": 384,
        "outputDim": 10,
        "dimensions": dimensions,
        "exportedAt": __import__("datetime").datetime.utcnow().isoformat() + "Z",
        "totalSamples": len(samples),
        "source": "bootstrap",
        "samples": samples,
    }

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(export_data, f, indent=2)

    print(f"Created bootstrap dataset with {len(samples)} samples: {output_path}")
    return output_path


if __name__ == "__main__":
    import sys

    output = sys.argv[1] if len(sys.argv) > 1 else "bootstrap_data.json"
    create_bootstrap_dataset(output)
