import yaml
from wand.image import Image
from wand.drawing import Drawing
from wand.color import Color
from typing import Dict, List, Any, Union

def parse_canvas(canvas_data: Union[List, Dict]) -> Dict:
    """Parse canvas in both full and minified formats"""
    if isinstance(canvas_data, list):
        # Minified format: [width, height, background]
        return {
            "width": canvas_data[0] if len(canvas_data) > 0 else 800,
            "height": canvas_data[1] if len(canvas_data) > 1 else 600,
            "background": canvas_data[2] if len(canvas_data) > 2 else "#ffffff"
        }
    elif isinstance(canvas_data, dict):
        # Full format - return as is
        return canvas_data
    return {"width": 800, "height": 600, "background": "#ffffff"}

def parse_style(style_data: Dict) -> Dict:
    """Parse style with both full and minified property names"""
    if not isinstance(style_data, dict):
        return {}
    
    parsed_style = {}
    
    # Handle both full and minified property names
    if "fill" in style_data:
        parsed_style["fill"] = style_data["fill"]
    elif "f" in style_data:
        parsed_style["fill"] = style_data["f"]
    
    if "stroke" in style_data:
        parsed_style["stroke"] = style_data["stroke"]
    elif "s" in style_data:
        parsed_style["stroke"] = style_data["s"]
    
    if "stroke_width" in style_data:
        parsed_style["stroke_width"] = style_data["stroke_width"]
    elif "sw" in style_data:
        parsed_style["stroke_width"] = style_data["sw"]
    
    if "opacity" in style_data:
        parsed_style["opacity"] = style_data["opacity"]
    elif "o" in style_data:
        parsed_style["opacity"] = style_data["o"]
    
    if "font_size" in style_data:
        parsed_style["font_size"] = style_data["font_size"]
    elif "fs" in style_data:
        parsed_style["font_size"] = style_data["fs"]
        
    if "font" in style_data:
        parsed_style["font"] = style_data["font"]
    elif "fn" in style_data:
        parsed_style["font"] = style_data["fn"]
    
    return parsed_style

def parse_shape(shape_data: Union[Dict, Any]) -> Dict:
    """Parse shape in both formats"""
    if not isinstance(shape_data, dict):
        return {}
    
    # Check if it's already in full format (has 'type' key)
    if "type" in shape_data:
        return shape_data
    
    # Handle minified format shapes
    parsed_shapes = []
    
    for key, value in shape_data.items():
        if key == "rectangle" and isinstance(value, list):
            # Minified: [[x, y], [width, height], "style", rx, ry, z]
            pos = value[0] if len(value) > 0 else [0, 0]
            size = value[1] if len(value) > 1 else [100, 100]
            style = value[2] if len(value) > 2 else None
            rx = value[3] if len(value) > 3 else 0
            ry = value[4] if len(value) > 4 else 0
            z = value[5] if len(value) > 5 else 0
            
            parsed_shapes.append({
                "type": "rectangle",
                "position": pos,
                "size": size,
                "style": style,
                "rx": rx,
                "ry": ry,
                "z": z
            })
            
        elif key == "circle" and isinstance(value, list):
            # Minified: [[cx, cy], radius, "style", z]
            center = value[0] if len(value) > 0 else [0, 0]
            radius = value[1] if len(value) > 1 else 10
            style = value[2] if len(value) > 2 else None
            z = value[3] if len(value) > 3 else 0
            
            parsed_shapes.append({
                "type": "circle",
                "center": center,
                "radius": radius,
                "style": style,
                "z": z
            })
            
        elif key == "ellipse" and isinstance(value, list):
            # Minified: [[cx, cy], [rx, ry], "style", z]
            center = value[0] if len(value) > 0 else [0, 0]
            radius = value[1] if len(value) > 1 else [10, 10]
            style = value[2] if len(value) > 2 else None
            z = value[3] if len(value) > 3 else 0
            
            parsed_shapes.append({
                "type": "ellipse",
                "center": center,
                "radius": radius,
                "style": style,
                "z": z
            })
            
        elif key == "line" and isinstance(value, list):
            # Minified: [[x1, y1], [x2, y2], "style", z]
            start = value[0] if len(value) > 0 else [0, 0]
            end = value[1] if len(value) > 1 else [100, 100]
            style = value[2] if len(value) > 2 else None
            z = value[3] if len(value) > 3 else 0
            
            parsed_shapes.append({
                "type": "line",
                "start": start,
                "end": end,
                "style": style,
                "z": z
            })
            
        elif key == "text" and isinstance(value, list):
            # Minified: ["content", [x, y], "fill", font_size, "font", "style", z]
            content = value[0] if len(value) > 0 else ""
            position = value[1] if len(value) > 1 else [0, 0]
            fill = value[2] if len(value) > 2 else "#000000"
            font_size = value[3] if len(value) > 3 else 20
            font = value[4] if len(value) > 4 else "Arial"
            style = value[5] if len(value) > 5 else None
            z = value[6] if len(value) > 6 else 0
            
            parsed_shapes.append({
                "type": "text",
                "content": content,
                "position": position,
                "fill": fill,
                "font_size": font_size,
                "font": font,
                "style": style,
                "z": z
            })
            
        elif key == "polygon" and isinstance(value, list):
            # Minified: [points_array, "style", z]
            points = value[0] if len(value) > 0 else []
            style = value[1] if len(value) > 1 else None
            z = value[2] if len(value) > 2 else 0
            
            parsed_shapes.append({
                "type": "polygon",
                "points": points,
                "style": style,
                "z": z
            })
            
        elif key == "polyline" and isinstance(value, list):
            # Minified: [points_array, "style", z]
            points = value[0] if len(value) > 0 else []
            style = value[1] if len(value) > 1 else None
            z = value[2] if len(value) > 2 else 0
            
            parsed_shapes.append({
                "type": "polyline",
                "points": points,
                "style": style,
                "z": z
            })
    
    return parsed_shapes[0] if len(parsed_shapes) == 1 else parsed_shapes

def apply_style(draw, shape, styles):
    style_name = shape.get("style")
    style = styles.get(style_name, {}) if style_name else {}
    
    fill = shape.get("fill", style.get("fill"))
    stroke = shape.get("stroke", style.get("stroke"))
    stroke_width = shape.get("stroke_width", style.get("stroke_width", 1))
    opacity = shape.get("opacity", style.get("opacity", 1.0))
    
    draw.fill_color = Color(fill) if fill else Color("transparent")
    draw.stroke_color = Color(stroke) if stroke else Color("transparent")
    draw.stroke_width = stroke_width
    draw.fill_opacity = opacity
    draw.stroke_opacity = opacity

def draw_shape(draw, img, shape, styles):
    t = shape["type"]
    draw_clone = Drawing()
    apply_style(draw_clone, shape, styles)
    
    if t == "rectangle":
        x, y = shape["position"]
        w, h = shape["size"]
        draw_clone.rectangle(left=x, top=y, width=w, height=h)
    elif t == "circle":
        cx, cy = shape["center"]
        r = shape["radius"]
        draw_clone.circle((cx, cy), (cx + r, cy))
    elif t == "ellipse":
        cx, cy = shape["center"]
        rx, ry = shape["radius"]
        draw_clone.ellipse(cx - rx, cy - ry, cx + rx, cy + ry)
    elif t == "line":
        sx, sy = shape["start"]
        ex, ey = shape["end"]
        draw_clone.line((sx, sy), (ex, ey))
    elif t == "polygon":
        pts = shape["points"]
        draw_clone.polygon(pts)
    elif t == "polyline":
        pts = shape["points"]
        draw_clone.polyline(pts)
    elif t == "text":
        x, y = shape["position"]
        draw_clone.font_size = shape.get("font_size", 20)
        draw_clone.fill_color = Color(shape.get("fill", "#000000"))
        draw_clone.font = shape.get("font", "Arial")
        draw_clone.text(x, y, shape["content"])
    else:
        print(f"⚠️ Unknown shape: {t}")
        return
    
    if shape.get("rotate"):
        angle = shape["rotate"]
        with img.clone() as temp_img:
            temp_img.background_color = Color("transparent")
            draw_clone.draw(temp_img)
            temp_img.rotate(angle)
            img.composite(temp_img, 0, 0)
    else:
        draw_clone(img)

def parse_yaml_config(config: Dict) -> Dict:
    """Parse the entire YAML config, handling both formats"""
    parsed_config = {}
    
    # Parse canvas
    parsed_config["canvas"] = parse_canvas(config.get("canvas", {}))
    
    # Parse styles
    styles = config.get("styles", {})
    parsed_styles = {}
    for name, style_data in styles.items():
        parsed_styles[name] = parse_style(style_data)
    parsed_config["styles"] = parsed_styles
    
    # Parse shapes
    shapes = config.get("shapes", [])
    parsed_shapes = []
    
    for shape_data in shapes:
        if isinstance(shape_data, dict):
            if "type" in shape_data:
                # Full format shape
                parsed_shapes.append(shape_data)
            else:
                # Minified format shape(s)
                parsed_shape = parse_shape(shape_data)
                if isinstance(parsed_shape, list):
                    parsed_shapes.extend(parsed_shape)
                else:
                    parsed_shapes.append(parsed_shape)
    
    parsed_config["shapes"] = parsed_shapes
    return parsed_config

def generate_from_yaml(yaml_path, output_file="output.png"):
    with open(yaml_path, "r") as f:
        config = yaml.safe_load(f)
    
    # Parse the config to handle both formats
    parsed_config = parse_yaml_config(config)
    
    canvas_cfg = parsed_config["canvas"]
    shapes = sorted(parsed_config.get("shapes", []), key=lambda x: x.get("z", 0))
    styles = parsed_config.get("styles", {})
    
    img = Image(width=canvas_cfg["width"], height=canvas_cfg["height"],
                background=Color(canvas_cfg["background"]))
    
    for shape in shapes:
        draw_shape(Drawing(), img, shape, styles)
    
    img.save(filename=output_file)
    print(f"✅ Image saved as {output_file}")

def generate_from_yaml_string(yaml_string: str, output_file="output.png"):
    """Generate image from YAML string instead of file"""
    config = yaml.safe_load(yaml_string)
    parsed_config = parse_yaml_config(config)
    
    canvas_cfg = parsed_config["canvas"]
    shapes = sorted(parsed_config.get("shapes", []), key=lambda x: x.get("z", 0))
    styles = parsed_config.get("styles", {})
    
    img = Image(width=canvas_cfg["width"], height=canvas_cfg["height"],
                background=Color(canvas_cfg["background"]))
    
    for shape in shapes:
        draw_shape(Drawing(), img, shape, styles)
    
    img.save(filename=output_file)
    print(f"✅ Image saved as {output_file}")

if __name__ == "__main__":

#unified
#     generate_from_yaml("demo.yaml")

#minified
     generate_from_yaml("minified.yaml")