import os
import random
import numpy as np
from PIL import Image, ImageDraw, ImageFilter

import json

# Define expanded classes (28 classes)
DEFAULT_CLASSES = [
    "apfel_elstar",
    "apfel_gala",
    "apfel_granny_smith",
    "banane_chiquita",
    "banane_bio",
    "birne",
    "orange",
    "zitrone",
    "limette",
    "erdbeere",
    "blaubeere",
    "weintraube_hell",
    "weintraube_dunkel",
    "pfirsich",
    "tomate",
    "gurke",
    "kartoffel",
    "karotte",
    "zwiebel_gelb",
    "zwiebel_rot",
    "knoblauch",
    "brokkoli",
    "paprika_rot",
    "paprika_gelb",
    "paprika_gruen",
    "champignon",
    "zucchini",
    "avocado"
]

def get_classes_from_db():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    articles_path = os.path.join(base_dir, "articles.json")
    if os.path.exists(articles_path):
        try:
            with open(articles_path, "r") as f:
                articles = json.load(f)
                keys = list(articles.keys())
                if keys:
                    return keys
        except Exception:
            pass
    return DEFAULT_CLASSES.copy()

CLASSES = get_classes_from_db()

def reload_classes():
    new_classes = get_classes_from_db()
    CLASSES.clear()
    CLASSES.extend(new_classes)
    print(f"Reloaded CLASSES: {len(CLASSES)} classes in memory.")

def draw_apple_elstar(draw):
    draw.ellipse([45, 50, 155, 160], fill=(210, 45, 45, 255))
    draw.ellipse([80, 45, 120, 60], fill=(210, 45, 45, 255))
    draw.ellipse([55, 70, 110, 140], fill=(230, 140, 40, 180))
    draw.line([(100, 55), (105, 25)], fill=(101, 67, 33, 255), width=4)
    draw.ellipse([93, 50, 107, 57], fill=(60, 20, 20, 255))

def draw_apple_gala(draw):
    draw.ellipse([45, 50, 155, 160], fill=(240, 190, 60, 255))
    draw.ellipse([80, 45, 120, 60], fill=(240, 190, 60, 255))
    for offset in range(-35, 45, 15):
        draw.arc([45 + offset, 50, 155 + offset, 160], start=30, end=150, fill=(210, 40, 40, 200), width=4)
        draw.arc([45 - offset, 50, 155 - offset, 160], start=30, end=150, fill=(210, 40, 40, 200), width=4)
    draw.line([(100, 55), (103, 25)], fill=(101, 67, 33, 255), width=4)
    draw.ellipse([93, 50, 107, 57], fill=(70, 50, 20, 255))

def draw_apple_granny_smith(draw):
    # Bright green apple
    draw.ellipse([45, 50, 155, 160], fill=(120, 200, 40, 255))
    draw.ellipse([80, 45, 120, 60], fill=(120, 200, 40, 255))
    # Light yellow highlight
    draw.ellipse([60, 70, 95, 110], fill=(160, 230, 80, 180))
    # Stem
    draw.line([(100, 55), (105, 25)], fill=(101, 67, 33, 255), width=4)
    draw.ellipse([93, 50, 107, 57], fill=(40, 70, 15, 255))

def draw_banana_chiquita(draw):
    banana_pts = [
        (40, 130), (50, 100), (70, 75), (100, 60), (130, 60),
        (160, 75), (170, 95), (160, 95), (135, 80), (110, 80),
        (85, 90), (65, 110), (50, 140)
    ]
    draw.polygon(banana_pts, fill=(245, 215, 50, 255))
    draw.polygon([(40, 130), (50, 140), (45, 145), (35, 135)], fill=(75, 100, 30, 255))
    draw.polygon([(160, 75), (170, 95), (175, 90), (165, 70)], fill=(90, 75, 30, 255))
    draw.ellipse([100, 65, 112, 75], fill=(10, 50, 150, 255))
    draw.ellipse([103, 68, 109, 72], fill=(245, 215, 50, 255))

def draw_banana_bio(draw):
    banana_pts = [
        (45, 125), (55, 100), (75, 80), (100, 70), (125, 70),
        (150, 82), (160, 98), (152, 98), (130, 88), (110, 85),
        (88, 93), (70, 110), (53, 133)
    ]
    draw.polygon(banana_pts, fill=(225, 220, 60, 255))
    draw.polygon([(45, 125), (53, 133), (48, 138), (40, 130)], fill=(50, 90, 20, 255))
    draw.polygon([(150, 82), (160, 98), (163, 93), (153, 77)], fill=(70, 90, 20, 255))
    spots = [(70, 105), (85, 92), (100, 83), (115, 83), (130, 86)]
    for sx, sy in spots:
        draw.ellipse([sx-2, sy-2, sx+2, sy+2], fill=(80, 50, 20, 200))
    draw.polygon([(95, 70), (105, 70), (102, 87), (92, 86)], fill=(120, 140, 90, 255))

def draw_pear(draw):
    # Birne: pear shape (narrow top, wide bottom)
    draw.ellipse([65, 50, 135, 110], fill=(180, 190, 50, 255)) # top
    draw.ellipse([50, 90, 150, 170], fill=(180, 190, 50, 255)) # bottom
    draw.ellipse([60, 75, 140, 145], fill=(180, 190, 50, 255)) # mid fill
    # Orange-red blush
    draw.ellipse([95, 100, 140, 150], fill=(210, 110, 30, 120))
    # Stem
    draw.line([(100, 55), (105, 25)], fill=(101, 67, 33, 255), width=3)

def draw_orange(draw):
    draw.ellipse([45, 45, 155, 155], fill=(245, 120, 20, 255))
    # Peel texture (little dots)
    for _ in range(15):
        tx = random.randint(60, 140)
        ty = random.randint(60, 140)
        draw.ellipse([tx, ty, tx+2, ty+2], fill=(215, 95, 5, 200))

def draw_lemon(draw):
    # Zitrone: yellow oval with pointed ends
    draw.ellipse([50, 65, 150, 135], fill=(245, 225, 40, 255))
    draw.polygon([(40, 100), (52, 90), (52, 110)], fill=(245, 225, 40, 255)) # left tip
    draw.polygon([(160, 100), (148, 90), (148, 110)], fill=(245, 225, 40, 255)) # right tip

def draw_lime(draw):
    # Limette: green oval with pointed ends
    draw.ellipse([55, 70, 145, 130], fill=(60, 160, 30, 255))
    draw.polygon([(47, 100), (57, 92), (57, 108)], fill=(60, 160, 30, 255))
    draw.polygon([(153, 100), (143, 92), (143, 108)], fill=(60, 160, 30, 255))

def draw_strawberry(draw):
    # Erdbeere: heart shape
    draw.polygon([(100, 165), (55, 80), (100, 60), (145, 80)], fill=(225, 30, 50, 255))
    draw.ellipse([55, 70, 105, 100], fill=(225, 30, 50, 255))
    draw.ellipse([95, 70, 145, 100], fill=(225, 30, 50, 255))
    # Seeds (yellow dots)
    for _ in range(20):
        sx = random.randint(65, 135)
        sy = random.randint(75, 140)
        draw.ellipse([sx, sy, sx+2, sy+2], fill=(240, 220, 100, 255))
    # Leafy top
    draw.polygon([(100, 65), (85, 45), (95, 60), (100, 40), (105, 60), (115, 45)], fill=(40, 160, 50, 255))

def draw_blueberry(draw):
    # Blaubeere: small blue circle
    draw.ellipse([70, 70, 130, 130], fill=(45, 80, 160, 255))
    # Crown (dark crown tip)
    draw.polygon([(90, 70), (100, 62), (110, 70), (105, 75), (95, 75)], fill=(20, 45, 95, 255))

def draw_grape_green(draw):
    # Weintraube hell: cluster of small green circles
    coords = [
        (100, 60), (85, 75), (115, 75), (70, 95), (100, 95), (130, 95),
        (85, 115), (115, 115), (100, 135), (100, 155)
    ]
    # Stem
    draw.line([(100, 60), (100, 35)], fill=(110, 140, 50, 255), width=3)
    for cx, cy in coords:
        draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=(160, 210, 80, 255))
        draw.ellipse([cx-16, cy-16, cx+16, cy+16], outline=(130, 180, 60, 255), width=1)

def draw_grape_dark(draw):
    # Weintraube dunkel: cluster of small dark purple circles
    coords = [
        (100, 60), (85, 75), (115, 75), (70, 95), (100, 95), (130, 95),
        (85, 115), (115, 115), (100, 135), (100, 155)
    ]
    draw.line([(100, 60), (100, 35)], fill=(90, 70, 110, 255), width=3)
    for cx, cy in coords:
        draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=(65, 30, 105, 255))
        draw.ellipse([cx-16, cy-16, cx+16, cy+16], outline=(45, 15, 80, 255), width=1)

def draw_peach(draw):
    # Pfirsich: fuzzy orange-pink circle
    draw.ellipse([45, 50, 155, 160], fill=(245, 130, 80, 255))
    draw.ellipse([47, 50, 120, 160], fill=(245, 80, 110, 255)) # pink cheek
    # Indentation line
    draw.arc([40, 50, 150, 160], start=270, end=90, fill=(185, 45, 45, 255), width=2)
    # stem
    draw.line([(100, 53), (102, 30)], fill=(90, 70, 40, 255), width=3)

def draw_tomato(draw):
    draw.ellipse([50, 55, 150, 155], fill=(225, 30, 30, 255))
    draw.ellipse([65, 70, 85, 85], fill=(255, 255, 255, 180))
    cx, cy = 100, 55
    stem_pts = [
        (cx, cy), (cx-15, cy-10), (cx-5, cy-2),
        (cx, cy-20), (cx+5, cy-2), (cx+15, cy-10),
        (cx+8, cy+5), (cx, cy+8), (cx-8, cy+5)
    ]
    draw.polygon(stem_pts, fill=(34, 139, 34, 255))
    draw.line([(cx, cy), (cx-5, cy-25)], fill=(34, 139, 34, 255), width=4)

def draw_cucumber(draw):
    draw.rounded_rectangle([35, 80, 165, 120], radius=20, fill=(20, 100, 40, 255))
    for x in range(50, 150, 15):
        draw.line([(x, 85), (x+5, 115)], fill=(30, 130, 55, 255), width=2)
    draw.ellipse([32, 95, 38, 105], fill=(139, 115, 85, 255))
    draw.ellipse([162, 95, 168, 105], fill=(100, 80, 40, 255))

def draw_potato(draw):
    # Kartoffel: irregular brown-yellow oval
    draw.rounded_rectangle([45, 65, 155, 135], radius=35, fill=(190, 160, 95, 255))
    # Add potato eyes (dark brown dots)
    eyes = [(65, 80), (135, 90), (115, 120), (80, 125), (100, 85)]
    for ex, ey in eyes:
        draw.ellipse([ex-2, ey-1, ex+2, ey+1], fill=(110, 85, 45, 255))
        draw.arc([ex-3, ey-3, ex+3, ey+3], start=0, end=180, fill=(110, 85, 45, 255), width=1)

def draw_carrot(draw):
    # Karotte: tapered orange shape
    draw.polygon([(40, 85), (160, 100), (40, 115)], fill=(245, 110, 20, 255))
    draw.ellipse([35, 85, 45, 115], fill=(245, 110, 20, 255))
    # Green leaves on one end (left)
    draw.polygon([(38, 100), (15, 80), (30, 95), (10, 100), (30, 105), (15, 120)], fill=(45, 155, 55, 255))
    # Texture lines
    for x in range(60, 140, 20):
        draw.line([(x, 92), (x, 108)], fill=(215, 85, 10, 255), width=2)

def draw_onion_yellow(draw):
    # Zwiebel gelb: golden bulb shape
    draw.ellipse([50, 60, 150, 150], fill=(195, 145, 75, 255))
    draw.polygon([(100, 40), (80, 65), (120, 65)], fill=(195, 145, 75, 255)) # top tip
    # Root hairs at bottom
    draw.line([(90, 150), (85, 162)], fill=(230, 210, 170, 255), width=2)
    draw.line([(100, 150), (100, 165)], fill=(230, 210, 170, 255), width=2)
    draw.line([(110, 150), (115, 162)], fill=(230, 210, 170, 255), width=2)

def draw_onion_red(draw):
    # Zwiebel rot: red-purple bulb shape
    draw.ellipse([50, 60, 150, 150], fill=(140, 35, 95, 255))
    draw.polygon([(100, 40), (80, 65), (120, 65)], fill=(140, 35, 95, 255))
    # Light stripes
    draw.arc([55, 60, 145, 150], start=180, end=360, fill=(185, 75, 135, 255), width=2)
    # Root hairs
    draw.line([(95, 150), (95, 162)], fill=(210, 190, 180, 255), width=2)
    draw.line([(105, 150), (105, 162)], fill=(210, 190, 180, 255), width=2)

def draw_garlic(draw):
    # Knoblauch: off-white bulb with segment lines
    draw.ellipse([50, 60, 150, 150], fill=(235, 230, 220, 255))
    draw.polygon([(100, 42), (85, 65), (115, 65)], fill=(235, 230, 220, 255))
    # Segment lines
    draw.arc([65, 60, 135, 150], start=270, end=90, fill=(195, 190, 180, 255), width=2)
    draw.arc([80, 60, 120, 150], start=270, end=90, fill=(195, 190, 180, 255), width=2)
    # Roots
    draw.line([(100, 150), (100, 160)], fill=(160, 150, 130, 255), width=2)

def draw_broccoli(draw):
    # Brokkoli: green tree-like structure
    # Stalk
    draw.rounded_rectangle([80, 110, 120, 170], radius=8, fill=(110, 175, 75, 255))
    # Florets (overlapping circles)
    draw.ellipse([60, 60, 110, 110], fill=(30, 110, 50, 255))
    draw.ellipse([90, 50, 150, 100], fill=(30, 110, 50, 255))
    draw.ellipse([75, 80, 135, 130], fill=(30, 110, 50, 255))
    # Add tiny dots/texture
    for _ in range(30):
        tx = random.randint(65, 135)
        ty = random.randint(55, 120)
        draw.ellipse([tx, ty, tx+2, ty+2], fill=(60, 150, 80, 255))

def draw_pepper_red(draw):
    # Paprika rot: blocky rounded shape
    draw.rounded_rectangle([50, 60, 150, 150], radius=25, fill=(210, 30, 30, 255))
    # Lobes indentations
    draw.arc([45, 60, 155, 150], start=230, end=310, fill=(160, 15, 15, 255), width=3)
    # Green stem
    draw.line([(100, 60), (100, 35)], fill=(45, 130, 45, 255), width=5)

def draw_pepper_yellow(draw):
    # Paprika gelb: blocky rounded shape
    draw.rounded_rectangle([50, 60, 150, 150], radius=25, fill=(240, 190, 20, 255))
    draw.arc([45, 60, 155, 150], start=230, end=310, fill=(195, 145, 10, 255), width=3)
    draw.line([(100, 60), (100, 35)], fill=(45, 130, 45, 255), width=5)

def draw_pepper_green(draw):
    # Paprika grün: blocky rounded shape
    draw.rounded_rectangle([50, 60, 150, 150], radius=25, fill=(35, 115, 45, 255))
    draw.arc([45, 60, 155, 150], start=230, end=310, fill=(20, 85, 25, 255), width=3)
    draw.line([(100, 60), (100, 35)], fill=(25, 80, 25, 255), width=5)

def draw_mushroom(draw):
    # Champignon: white-brown stalk and cap
    # Stalk
    draw.rounded_rectangle([85, 100, 115, 160], radius=6, fill=(225, 215, 205, 255))
    # Cap (umbrella)
    draw.ellipse([45, 60, 155, 120], fill=(210, 200, 190, 255))
    draw.rectangle([45, 90, 155, 120], fill=(210, 200, 190, 255)) # bottom flat
    # Gills (brown under cap)
    draw.line([(45, 120), (155, 120)], fill=(120, 100, 90, 255), width=3)

def draw_zucchini(draw):
    # Zucchini: long cylinder, thicker than cucumber, dark green
    draw.rounded_rectangle([30, 75, 170, 125], radius=22, fill=(15, 75, 35, 255))
    # Speckled light green stripes
    for x in range(45, 155, 20):
        draw.line([(x, 80), (x+8, 120)], fill=(35, 115, 60, 180), width=3)
    # Thick stem
    draw.ellipse([167, 95, 173, 105], fill=(80, 95, 50, 255))

def draw_avocado(draw):
    # Halved avocado: pear outline, yellow-green center, brown seed
    # Outer skin
    draw.ellipse([60, 50, 140, 120], fill=(15, 45, 20, 255))
    draw.ellipse([45, 85, 155, 170], fill=(15, 45, 20, 255))
    # Light green inner meat
    draw.ellipse([63, 53, 137, 117], fill=(185, 215, 110, 255))
    draw.ellipse([49, 89, 151, 166], fill=(185, 215, 110, 255))
    # Yellow center
    draw.ellipse([58, 98, 142, 158], fill=(215, 235, 135, 255))
    # Brown pit (seed) in center
    draw.ellipse([78, 105, 122, 149], fill=(110, 65, 30, 255))
    draw.ellipse([82, 109, 100, 125], fill=(255, 255, 255, 80)) # Pit reflection

def create_object_image(class_idx, angle, scale):
    obj_img = Image.new("RGBA", (200, 200), (0, 0, 0, 0))
    draw = ImageDraw.Draw(obj_img)
    
    # Map class index to drawing functions
    draw_funcs = [
        draw_apple_elstar, draw_apple_gala, draw_apple_granny_smith,
        draw_banana_chiquita, draw_banana_bio, draw_pear,
        draw_orange, draw_lemon, draw_lime,
        draw_strawberry, draw_blueberry, draw_grape_green,
        draw_grape_dark, draw_peach, draw_tomato,
        draw_cucumber, draw_potato, draw_carrot,
        draw_onion_yellow, draw_onion_red, draw_garlic,
        draw_broccoli, draw_pepper_red, draw_pepper_yellow,
        draw_pepper_green, draw_mushroom, draw_zucchini,
        draw_avocado
    ]
    
    if class_idx < len(draw_funcs):
        draw_funcs[class_idx](draw)
        
    rotated = obj_img.rotate(angle, resample=Image.Resampling.BILINEAR, expand=True)
    
    if scale != 1.0:
        new_w = int(rotated.width * scale)
        new_h = int(rotated.height * scale)
        rotated = rotated.resize((new_w, new_h), resample=Image.Resampling.BILINEAR)
        
    bbox = rotated.getbbox()
    if bbox is None:
        return None, None
        
    cropped = rotated.crop(bbox)
    return cropped, bbox

def generate_background(width, height):
    bg = Image.new("RGB", (width, height), (210, 215, 220))
    draw = ImageDraw.Draw(bg)
    grid_color = (185, 190, 195)
    for x in range(0, width, 80):
        draw.line([(x, 0), (x, height)], fill=grid_color, width=2)
    for y in range(0, height, 80):
        draw.line([(0, y), (width, y)], fill=grid_color, width=2)
    draw.rectangle([10, 10, width-10, height-10], outline=(170, 175, 180), width=3)
    
    arr = np.array(bg, dtype=np.float32)
    x_grad = np.linspace(-15, 15, width)
    y_grad = np.linspace(-15, 15, height)
    xx, yy = np.meshgrid(x_grad, y_grad)
    arr[:, :, 0] += xx + yy
    arr[:, :, 1] += xx + yy
    arr[:, :, 2] += xx + yy
    noise = np.random.normal(0, 3, (height, width, 3))
    arr += noise
    arr = np.clip(arr, 0, 255).astype(np.uint8)
    return Image.fromarray(arr)

def generate_sample(width=640, height=480):
    bg = generate_background(width, height)
    num_objects = random.randint(1, 3)
    annotations = []
    placed_boxes = []
    
    for _ in range(num_objects):
        class_idx = random.randint(0, len(CLASSES)-1)
        angle = random.randint(0, 360)
        scale = random.uniform(0.7, 1.2)
        
        obj_img, _ = create_object_image(class_idx, angle, scale)
        if obj_img is None:
            continue
            
        ow, oh = obj_img.size
        placed = False
        
        for attempt in range(25):
            x = random.randint(30, width - ow - 30)
            y = random.randint(30, height - oh - 30)
            
            overlap = False
            for px1, py1, px2, py2 in placed_boxes:
                ix1 = max(x, px1)
                iy1 = max(y, py1)
                ix2 = min(x + ow, px2)
                iy2 = min(y + oh, py2)
                
                if ix2 > ix1 and iy2 > iy1:
                    inter_area = (ix2 - ix1) * (iy2 - iy1)
                    box_area = ow * oh
                    p_area = (px2 - px1) * (py2 - py1)
                    iou = inter_area / float(box_area + p_area - inter_area)
                    if iou > 0.12:
                        overlap = True
                        break
            
            if not overlap:
                bg.paste(obj_img, (x, y), obj_img)
                placed_boxes.append((x, y, x + ow, y + oh))
                x_center = (x + ow / 2.0) / width
                y_center = (y + oh / 2.0) / height
                w_norm = ow / float(width)
                h_norm = oh / float(height)
                annotations.append(f"{class_idx} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
                placed = True
                break
                
        if not placed and len(placed_boxes) == 0:
            x = random.randint(30, width - ow - 30)
            y = random.randint(30, height - oh - 30)
            bg.paste(obj_img, (x, y), obj_img)
            placed_boxes.append((x, y, x + ow, y + oh))
            x_center = (x + ow / 2.0) / width
            y_center = (y + oh / 2.0) / height
            w_norm = ow / float(width)
            h_norm = oh / float(height)
            annotations.append(f"{class_idx} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
            
    bg = bg.filter(ImageFilter.GaussianBlur(0.5))
    return bg, annotations

def build_dataset(base_dir="dataset", train_count=200, val_count=50):
    os.makedirs(os.path.join(base_dir, "images", "train"), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "images", "val"), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "labels", "train"), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "labels", "val"), exist_ok=True)
    
    with open(os.path.join(base_dir, "classes.txt"), "w") as f:
        f.write("\n".join(CLASSES))
        
    print(f"Generating {train_count} training samples for {len(CLASSES)} classes...")
    for idx in range(train_count):
        img, ann = generate_sample()
        img.save(os.path.join(base_dir, "images", "train", f"train_{idx:04d}.jpg"), quality=90)
        with open(os.path.join(base_dir, "labels", "train", f"train_{idx:04d}.txt"), "w") as f:
            f.write("\n".join(ann))
            
    print(f"Generating {val_count} validation samples...")
    for idx in range(val_count):
        img, ann = generate_sample()
        img.save(os.path.join(base_dir, "images", "val", f"val_{idx:04d}.jpg"), quality=90)
        with open(os.path.join(base_dir, "labels", "val", f"val_{idx:04d}.txt"), "w") as f:
            f.write("\n".join(ann))
            
    print("Dataset generation complete!")

if __name__ == "__main__":
    build_dataset(train_count=200, val_count=50)