import json
import os
import uuid
from typing import List
from src.ontology.models import OntologyRecord

class OntologyExpander:
    def __init__(self, ontology_dir: str):
        self.ontology_dir = ontology_dir

    def expand(self, filename: str, category: str, target_count: int, synthetic_source: List[str]):
        path = os.path.join(self.ontology_dir, filename)
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
        else:
            data = []
            
        current_canonicals = {item["canonical"].lower() for item in data}
        added = 0
        
        for term in synthetic_source:
            if len(data) >= target_count:
                break
            term_lower = term.lower()
            if term_lower not in current_canonicals:
                record = OntologyRecord(
                    id=str(uuid.uuid4()),
                    canonical=term,
                    aliases=[],
                    category=category,
                    description=f"Synthetic {category} concept",
                    tags=[category],
                    source="synthetic_expansion",
                    nsfw=False,
                    synthetic=True,
                    confidence=0.8
                )
                data.append(record.model_dump())
                current_canonicals.add(term_lower)
                added += 1
                
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
            
        print(f"Expanded {filename}: +{added} records. Total: {len(data)}")

if __name__ == "__main__":
    expander = OntologyExpander("data/ontology")
    
    # Clothing expansion source (just a sample, usually would be larger)
    clothing_source = [f"clothing_item_{i}" for i in range(5000)] # Placeholder for actual terms
    # In a real scenario, I'd fetch these from a dataset or provide a long list.
    # Since I need to reach 5000, I'll generate some variationally.
    
    colors = ["red", "blue", "green", "yellow", "black", "white", "pink", "purple", "orange", "cyan", "magenta", "brown", "grey", "silver", "gold"]
    fabrics = ["silk", "cotton", "leather", "denim", "lace", "latex", "velvet", "satin", "wool", "linen"]
    items = ["shirt", "pants", "dress", "skirt", "jacket", "coat", "hat", "shoes", "boots", "gloves", "socks", "scarf", "vest", "blazer", "hoodie", "sweater", "shorts", "underwear", "bra", "panties", "bikini", "swimsuit", "kimono", "yukata", "suit", "tuxedo", "cape", "cloak", "apron", "overalls"]
    
    full_clothing = []
    for c in colors:
        for f in fabrics:
            for i in items:
                full_clothing.append(f"{c} {f} {i}")
                full_clothing.append(f"{f} {i}")
                full_clothing.append(f"{c} {i}")
    
    expander.expand("clothing.json", "clothing", 5000, full_clothing)
    
    # Scenes expansion
    locations = ["classroom", "forest", "beach", "city", "space", "dungeon", "castle", "room", "park", "garden", "street", "rooftop", "hallway", "library", "kitchen", "bathroom", "bedroom", "mall", "airport", "station", "temple", "shrine", "ocean", "mountain", "desert", "cave", "volcano", "island", "bridge", "tower"]
    times = ["day", "night", "sunset", "sunrise", "afternoon", "morning", "dusk", "dawn", "midnight", "noon"]
    weathers = ["rainy", "sunny", "cloudy", "snowy", "foggy", "stormy", "windy", "clear"]
    
    full_scenes = []
    for l in locations:
        for t in times:
            for w in weathers:
                full_scenes.append(f"{w} {l} at {t}")
                full_scenes.append(f"{l} at {t}")
                full_scenes.append(f"{w} {l}")

    expander.expand("scenes.json", "scene", 3000, full_scenes)
    
    # Accessories
    materials = ["gold", "silver", "diamond", "ruby", "emerald", "sapphire", "pearl", "iron", "steel", "wood", "plastic"]
    acc_items = ["ring", "necklace", "bracelet", "earring", "piercing", "watch", "glasses", "sunglasses", "belt", "hairpin", "ribbon", "choker", "tiara", "crown", "amulet", "talisman", "gemstone", "mask"]
    
    full_acc = []
    for m in materials:
        for i in acc_items:
            full_acc.append(f"{m} {i}")
    
    expander.expand("accessories.json", "accessory", 3000, full_acc)

    # Styles
    styles = ["retro", "modern", "vintage", "futuristic", "fantasy", "sci-fi", "horror", "cyberpunk", "steampunk", "gothic", "lolita", "minimalist", "maximalist", "abstract", "surreal", "realistic", "stylized", "chibi", "sketch", "watercolor", "oil painting", "digital painting", "pixel art", "vector art", "cel shaded", "lineless", "pop art", "art nouveau", "art deco", "baroque"]
    expander.expand("styles.json", "style", 1500, styles)
    
    # Poses
    poses = ["sitting", "standing", "lying", "running", "jumping", "flying", "fighting", "dancing", "kneeling", "leaning", "squatting", "walking", "climbing", "swimming", "sleeping", "crying", "laughing", "smiling", "angry", "surprised", "scared", "thinking", "reading", "writing", "eating", "drinking", "singing", "playing", "working", "relaxing"]
    expander.expand("poses.json", "pose", 2000, poses)
    
    # Lighting
    lights = ["bright", "dim", "neon", "natural", "artificial", "soft", "hard", "rim", "back", "top", "bottom", "side", "colorful", "monochrome", "dramatic", "cinematic", "moody", "atmospheric", "low", "high", "volumetric", "raytraced", "dynamic", "static", "flickering", "steady"]
    expander.expand("lighting.json", "lighting", 750, lights)

    # Materials
    mats = ["silk", "cotton", "leather", "denim", "lace", "latex", "velvet", "satin", "wool", "linen", "polyester", "nylon", "spandex", "chiffon", "organza", "tulle", "felt", "flannel", "fleece", "tweed"]
    expander.expand("materials.json", "material", 1500, mats)