import argparse
import json
import random


TEMPLATES = [
    {
        "instruction": "Fix the Python code",
        "input": "def add(a,b) return a+b",
        "output": "def add(a, b): return a + b",
        "explanation": "Added missing colon and corrected syntax.",
    },
    {
        "instruction": "Fix loop syntax",
        "input": "for i in range(5 print(i)",
        "output": "for i in range(5): print(i)",
        "explanation": "Added missing parenthesis and colon.",
    },
    {
        "instruction": "Fix condition",
        "input": "if x = 10: print(x)",
        "output": "if x == 10: print(x)",
        "explanation": "Corrected assignment to comparison operator.",
    },
    {
        "instruction": "Explain code",
        "input": "for i in range(3): print(i)",
        "output": "Prints numbers from 0 to 2.",
        "explanation": "Loop iterates from 0 to 2 and prints values.",
    },
    {
        "instruction": "Write a Python function",
        "input": "Create a function to multiply two numbers",
        "output": "def multiply(a, b):\n    return a * b",
        "explanation": "Defined a multiply function that returns the product of two inputs.",
    },
    {
        "instruction": "Write a Python function",
        "input": "Create a function to add two numbers",
        "output": "def add(a, b):\n    return a + b",
        "explanation": "Defined an add function that returns the sum of two inputs.",
    },
    {
        "instruction": "Write a Python function",
        "input": "Create a function to subtract two numbers",
        "output": "def subtract(a, b):\n    return a - b",
        "explanation": "Defined a subtract function that returns the difference between two inputs.",
    },
    {
        "instruction": "Write a Python function",
        "input": "Create a function to divide two numbers",
        "output": "def divide(a, b):\n    return a / b",
        "explanation": "Defined a divide function that returns the quotient of two inputs.",
    },
]


def format_training_text(template):
    target = {
        "code": template["output"],
        "explanation": template["explanation"],
    }
    return (
        f"Instruction: {template['instruction']}\n"
        f"Input: {template['input']}\n"
        "Return only valid JSON with keys code and explanation.\n"
        f"JSON: {json.dumps(target, ensure_ascii=False)}\n"
    )


def generate_sample():
    template = random.choice(TEMPLATES)
    text = format_training_text(template)
    return {
        "instruction": template["instruction"],
        "input": template["input"],
        "output": template["output"],
        "explanation": template["explanation"],
        "text": text,
        "confidence": round(random.uniform(0.9, 0.99), 2),
        "relevancy": round(random.uniform(0.85, 0.99), 2),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--size", type=int, default=8000)
    parser.add_argument("--out", type=str, default="train.json")
    args = parser.parse_args()

    if not (5000 <= args.size <= 10000):
        raise ValueError("size must be between 5000 and 10000")

    dataset = [generate_sample() for _ in range(args.size)]
    with open(args.out, "w", encoding="utf-8") as f:
        json.dump(dataset, f, indent=2)
    print(f"Dataset created: {len(dataset)} -> {args.out}")


if __name__ == "__main__":
    main()