src/skill_generator.py

"""
Skill Generator — Free dataset generation for new skills.
Uses HuggingFace datasets for high-quality training data.
No API key needed.
"""

import json
import random
from pathlib import Path
from typing import List, Dict, Optional
from dataclasses import dataclass

try:
    from datasets import load_dataset
    HAS_DATASETS = True
except ImportError:
    HAS_DATASETS = False


@dataclass
class SkillTemplate:
    name: str
    description: str
    token: str
    trigger_patterns: List[str]
    system_prompt: str
    question_templates: List[str]
    num_examples: int = 200

    def generate_examples(self) -> List[Dict[str, str]]:
        examples = []
        
        # Try to use HuggingFace datasets for higher quality
        if HAS_DATASETS:
            examples = self._generate_from_hf()
        
        # Fallback to template-based if HF fails
        if not examples:
            examples = self._generate_from_templates()
        
        return examples

    def _generate_from_hf(self) -> List[Dict[str, str]]:
        """Generate examples from HuggingFace datasets"""
        examples = []
        try:
            if self.name == "code_expert":
                ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train", streaming=True)
                for i, row in enumerate(ds):
                    if i >= self.num_examples:
                        break
                    examples.append({
                        "prompt": row["instruction"],
                        "response": row["output"],
                        "skill_token": self.token,
                        "system_prompt": self.system_prompt
                    })
            elif self.name == "math_solver":
                ds = load_dataset("openai/gsm8k", "main", split="train", streaming=True)
                for i, row in enumerate(ds):
                    if i >= self.num_examples:
                        break
                    examples.append({
                        "prompt": row["question"],
                        "response": row["answer"],
                        "skill_token": self.token,
                        "system_prompt": self.system_prompt
                    })
            elif self.name == "creative_writer":
                ds = load_dataset("HuggingFaceH4/ultrachat_200k", "default", split="train_sft", streaming=True)
                for i, row in enumerate(ds):
                    if i >= self.num_examples * 3:
                        break
                    # Extract first user/assistant pair
                    if row.get("messages"):
                        msgs = row["messages"]
                        for j in range(len(msgs) - 1):
                            if msgs[j].get("role") == "user" and msgs[j + 1].get("role") == "assistant":
                                examples.append({
                                    "prompt": msgs[j]["content"],
                                    "response": msgs[j + 1]["content"],
                                    "skill_token": self.token,
                                    "system_prompt": self.system_prompt
                                })
                                break
                        if len(examples) >= self.num_examples:
                            break
            elif self.name == "data_analyst":
                ds = load_dataset("HuggingFaceH4/ultrachat_200k", "default", split="train_sft", streaming=True)
                for i, row in enumerate(ds):
                    if i >= self.num_examples * 10:
                        break
                    if row.get("messages"):
                        msgs = row["messages"]
                        # find first user->assistant pair mentioning data topics
                        for j in range(len(msgs) - 1):
                            if msgs[j].get("role") == "user" and msgs[j + 1].get("role") == "assistant":
                                if any(w in msgs[j]["content"].lower() for w in ["data", "analyze", "chart", "statistics", "dataset", "visualization"]):
                                    examples.append({
                                        "prompt": msgs[j]["content"],
                                        "response": msgs[j + 1]["content"],
                                        "skill_token": self.token,
                                        "system_prompt": self.system_prompt
                                    })
                                    break
                            if len(examples) >= self.num_examples:
                                break
                    if len(examples) >= self.num_examples:
                        break
            elif self.name == "translator":
                ds = load_dataset("Helsinki-NLP/opus-100", "en-fr", split="train", streaming=True)
                for i, row in enumerate(ds):
                    if i >= self.num_examples:
                        break
                    tr = row.get("translation", {})
                    en, fr = tr.get("en", ""), tr.get("fr", "")
                    if not en or not fr:
                        continue
                    examples.append({
                        "prompt": f"Translate to French: {en}",
                        "response": fr,
                        "skill_token": self.token,
                        "system_prompt": self.system_prompt
                    })
            elif self.name == "reasoning":
                ds = load_dataset("openai/gsm8k", "main", split="train", streaming=True)
                for i, row in enumerate(ds):
                    if i >= self.num_examples:
                        break
                    examples.append({
                        "prompt": f"Solve step by step: {row['question']}",
                        "response": row["answer"],
                        "skill_token": self.token,
                        "system_prompt": self.system_prompt
                    })
        except Exception as e:
            print(f"Warning: Could not load HF dataset for {self.name}: {e}")
            examples = []
        
        return examples

    def _generate_from_templates(self) -> List[Dict[str, str]]:
        """Fallback template-based generation"""
        placeholders = {
            'action': ['sort a list', 'reverse a string', 'find duplicates', 'validate email', 'parse JSON', 'merge dictionaries'],
            'code_snippet': ['def foo(): pass', 'x = [1,2,3]', 'for i in range(10): print(i)'],
            'concept': ['recursion', 'closures', 'decorators', 'generators', 'async/await', 'OOP'],
            'framework': ['Flask', 'FastAPI', 'Django', 'React', 'pandas', 'PyTorch'],
            'algorithm': ['binary search', 'quicksort', 'merge sort', 'BFS', 'DFS', 'dynamic programming'],
            'function': ['sin(x)', 'x^2 + 2x + 1', 'e^x', '1/x', 'log(x)'],
            'equation': ['2x + 5 = 15', 'x^2 - 4 = 0', '3x + 2y = 12'],
            'theorem': ['Pythagorean theorem', 'binomial theorem', 'intermediate value theorem'],
            'system_eq': ['x + y = 10, x - y = 4', '2x + y = 7, x - 3y = -5'],
            'polynomial': ['x^2 - 5x + 6', 'x^3 - 2x^2 - x + 2'],
            'topic': ['space exploration', 'artificial intelligence', 'climate change', 'technology', 'nature'],
            'genre': ['science fiction', 'mystery', 'fantasy', 'horror', 'thriller'],
            'setting': ['Mars colony', 'medieval kingdom', 'underwater city', 'parallel universe'],
            'characters': ['a robot and a human', 'time travelers', 'detective and suspect'],
            'scene': ['a bustling marketplace', 'an abandoned spaceship', 'a magical forest'],
            'character_type': ['anti-hero', 'reluctant mentor', 'mad scientist'],
            'dataset_desc': ['sales data for Q1-Q4', 'customer survey responses', 'website traffic logs'],
            'data': ['monthly revenue', 'user engagement metrics', 'weather data', 'stock prices'],
            'data_type': ['time series', 'categorical', 'geospatial'],
            'ml_problem': ['customer churn', 'image classification', 'sentiment analysis'],
            'language': ['Spanish', 'French', 'German', 'Japanese', 'Chinese'],
            'text': ['Hello, how are you?', 'The weather is nice', 'I love programming'],
            'phrase': ['good morning', 'how much', 'where is', 'nice to meet you'],
            'puzzle': ['Three switches control three bulbs', 'You have 8 balls, one heavier'],
            'premises': ['all humans are mortal', 'Socrates is human', 'All birds can fly'],
            'riddle': ['What has keys but no locks?', 'I speak without a mouth'],
            'sequence': ['2, 4, 8, 16, ?', '1, 1, 2, 3, 5, ?'],
        }

        examples = []
        for i in range(self.num_examples):
            q_template = random.choice(self.question_templates)
            params = {k: random.choice(v) for k, v in placeholders.items()}
            question = q_template.format(**params)
            examples.append({
                "prompt": question,
                "skill_token": self.token,
                "system_prompt": self.system_prompt
            })
        return examples

    def save_dataset(self, output_path: str):
        examples = self.generate_examples()
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        # One JSON object per line (JSONL). load_jsonl in train.py / the
        # notebook parses each line back into a dict, so a multi-line ChatML
        # blob would get fragmented into one broken example per line.
        with open(path, 'w', encoding='utf-8') as f:
            for ex in examples:
                row = {
                    "prompt": ex.get("prompt", ex.get("question", "")),
                    "response": ex.get("response", "Here is a helpful response."),
                    "skill_token": ex.get("skill_token", ""),
                    "system_prompt": ex.get("system_prompt", ""),
                }
                f.write(json.dumps(row, ensure_ascii=False) + "\n")

        print(f"Generated {len(examples)} examples -> {path}")
        return examples

    def save_skill_file(self, output_path: str):
        skill_data = {
            "name": self.name,
            "token": self.token,
            "description": self.description,
            "trigger_patterns": self.trigger_patterns,
            "system_prompt": self.system_prompt,
            "num_examples": self.num_examples
        }
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            json.dump(skill_data, f, indent=2)
        print(f"Skill template saved -> {path}")


SKILL_TEMPLATES = {
    "code_expert": SkillTemplate(
        name="code_expert",
        description="Expert programmer - writes, debugs, and explains code",
        token="<SKILL:code>",
        trigger_patterns=["code", "python", "function", "debug", "program", "script", "algorithm", "api"],
        system_prompt="You are an expert programmer. Write clean, efficient, well-documented code. Always explain your approach.",
        question_templates=[
            "Write a Python function that {action}",
            "Create a {action} in Python",
            "Debug this code: {code_snippet}",
            "Explain how {concept} works in programming",
            "Write a {action} using {framework}",
            "Optimize this function for performance: {code_snippet}",
            "Implement {algorithm} in Python",
            "Create a REST API endpoint for {action}",
        ],
        num_examples=200
    ),

    "math_solver": SkillTemplate(
        name="math_solver",
        description="Advanced mathematics - solves equations, proofs, and problems step by step",
        token="<SKILL:math>",
        trigger_patterns=["math", "equation", "solve", "calculate", "proof", "theorem", "integral", "derivative", "algebra", "calculus"],
        system_prompt="You are a mathematics expert. Show all steps clearly. Verify your answers.",
        question_templates=[
            "Solve for x: {equation}",
            "Find the derivative of {function}",
            "Calculate the integral of {function}",
            "Prove that {theorem}",
            "Solve this system of equations: {system_eq}",
            "Find the limit as x approaches a value",
            "Factorize {polynomial}",
            "Solve the differential equation: {equation}",
        ],
        num_examples=200
    ),

    "creative_writer": SkillTemplate(
        name="creative_writer",
        description="Creative writing - stories, poems, essays, and scripts",
        token="<SKILL:write>",
        trigger_patterns=["write", "story", "poem", "essay", "creative", "script", "narrative", "fiction"],
        system_prompt="You are a creative writer. Be imaginative, vivid, and engaging. Use strong imagery and varied sentence structure.",
        question_templates=[
            "Write a short story about {topic}",
            "Compose a poem about {topic}",
            "Write an essay on {topic}",
            "Create a dialogue between {characters}",
            "Write a {genre} story set in {setting}",
            "Describe {scene} in vivid detail",
            "Write a sonnet about {topic}",
            "Create a character description for a {character_type}",
        ],
        num_examples=150
    ),

    "data_analyst": SkillTemplate(
        name="data_analyst",
        description="Data analysis - interprets data, creates insights, suggests visualizations",
        token="<SKILL:data>",
        trigger_patterns=["data", "analyze", "statistics", "chart", "graph", "dataset", "pandas", "visualization"],
        system_prompt="You are a data analyst. Be precise with numbers. Suggest appropriate visualizations. Explain your methodology.",
        question_templates=[
            "Analyze this dataset: {dataset_desc}",
            "What insights can you find in this data: {data}",
            "Create a visualization plan for {data_type}",
            "Calculate statistics for: {data}",
            "What trends do you see in {data}",
            "Suggest a machine learning approach for {ml_problem}",
            "Clean and preprocess this data: {dataset_desc}",
        ],
        num_examples=150
    ),

    "translator": SkillTemplate(
        name="translator",
        description="Multi-language translator - accurate, context-aware translation",
        token="<SKILL:translate>",
        trigger_patterns=["translate", "translation", "spanish", "french", "german", "chinese", "japanese", "language"],
        system_prompt="You are a professional translator. Preserve tone, context, and cultural nuances. Provide both translation and explanation.",
        question_templates=[
            "Translate to {language}: {text}",
            "How do you say {phrase} in {language}?",
            "Translate this {language} text to English: {text}",
            "What's the {language} equivalent of {phrase}?",
            "Translate and explain the cultural context: {text}",
        ],
        num_examples=200
    ),

    "reasoning": SkillTemplate(
        name="reasoning",
        description="Logical reasoning - solves puzzles, logic problems, and analytical questions",
        token="<SKILL:logic>",
        trigger_patterns=["logic", "puzzle", "riddle", "reason", "think", "analyze", "deduce", "infer"],
        system_prompt="You are a logical reasoning expert. Break problems into steps. Consider all possibilities before concluding.",
        question_templates=[
            "Solve this logic puzzle: {puzzle}",
            "If {premises}, what can we conclude?",
            "Deduce the answer: {puzzle}",
            "Solve this riddle: {riddle}",
            "What's the pattern in: {sequence}",
            "Reason through this problem: {puzzle}",
            "If all A are B, and some B are C, then what follows?",
        ],
        num_examples=150
    )
}


def generate_all_skills(output_dir: str = "skills"):
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    for skill_name, template in SKILL_TEMPLATES.items():
        dataset_path = output_path.parent / "datasets" / f"{skill_name}_dataset.jsonl"
        skill_path = output_path / f"{skill_name}.skill"

        template.save_dataset(str(dataset_path))
        template.save_skill_file(str(skill_path))

    print(f"\nGenerated {len(SKILL_TEMPLATES)} skills in {output_path}")


def generate_custom_skill(
    name: str,
    description: str,
    trigger_patterns: List[str],
    system_prompt: str,
    num_examples: int = 100,
    output_dir: str = "skills"
):
    token = f"<SKILL:{name}>"
    template = SkillTemplate(
        name=name,
        description=description,
        token=token,
        trigger_patterns=trigger_patterns,
        system_prompt=system_prompt,
        question_templates=["{question}"],
        num_examples=num_examples
    )

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    dataset_path = output_path.parent / "datasets" / f"{name}_dataset.jsonl"
    skill_path = output_path / f"{name}.skill"

    template.save_dataset(str(dataset_path))
    template.save_skill_file(str(skill_path))

    print(f"Custom skill '{name}' generated")
    return template


if __name__ == "__main__":
    generate_all_skills("skills")