schema_visualizer.py

"""
Database Schema Visualization Module
Creates ERD-style diagrams similar to dbdiagram.io
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import json
from typing import Dict, List, Tuple, Optional
import math


class SchemaVisualizer:
    """Creates interactive database schema visualizations"""
    
    def __init__(self, parent, theme_manager):
        self.parent = parent
        self.theme = theme_manager
        self.tables = {}
        self.relationships = []
        
        # Visualization settings
        self.table_width = 200
        self.table_header_height = 35
        self.field_height = 25
        self.table_padding = 50
        
        # Canvas state
        self.zoom_level = 1.0
        self.pan_offset = [0, 0]
        self.dragging = False
        self.drag_start = None
        
        self._create_ui()
    
    def _create_ui(self):
        """Create visualization UI"""
        # Control frame
        control_frame = ttk.Frame(self.parent)
        control_frame.pack(side=tk.TOP, fill=tk.X, padx=5, pady=5)
        
        ttk.Button(control_frame, text="🔍 Zoom In", 
                  command=self.zoom_in).pack(side=tk.LEFT, padx=2)
        ttk.Button(control_frame, text="🔍 Zoom Out", 
                  command=self.zoom_out).pack(side=tk.LEFT, padx=2)
        ttk.Button(control_frame, text="↺ Reset View", 
                  command=self.reset_view).pack(side=tk.LEFT, padx=2)
        ttk.Button(control_frame, text="💾 Export PNG", 
                  command=self.export_image).pack(side=tk.LEFT, padx=2)
        
        ttk.Separator(control_frame, orient=tk.VERTICAL).pack(side=tk.LEFT, 
                                                              padx=5, fill=tk.Y)
        
        ttk.Label(control_frame, text="Layout:").pack(side=tk.LEFT, padx=5)
        self.layout_var = tk.StringVar(value="auto")
        layout_combo = ttk.Combobox(control_frame, textvariable=self.layout_var,
                                    values=["auto", "grid", "circular"],
                                    state="readonly", width=10)
        layout_combo.pack(side=tk.LEFT, padx=2)
        layout_combo.bind("<<ComboboxSelected>>", lambda e: self.redraw())
        
        # Canvas with scrollbars
        canvas_frame = ttk.Frame(self.parent)
        canvas_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # Scrollbars
        h_scroll = ttk.Scrollbar(canvas_frame, orient=tk.HORIZONTAL)
        h_scroll.pack(side=tk.BOTTOM, fill=tk.X)
        
        v_scroll = ttk.Scrollbar(canvas_frame, orient=tk.VERTICAL)
        v_scroll.pack(side=tk.RIGHT, fill=tk.Y)
        
        # Canvas
        self.canvas = tk.Canvas(
            canvas_frame,
            bg=self.theme.get_color("bg"),
            highlightthickness=0,
            xscrollcommand=h_scroll.set,
            yscrollcommand=v_scroll.set
        )
        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        h_scroll.config(command=self.canvas.xview)
        v_scroll.config(command=self.canvas.yview)
        
        # Bind events
        self.canvas.bind("<MouseWheel>", self._on_mousewheel)
        self.canvas.bind("<ButtonPress-1>", self._on_drag_start)
        self.canvas.bind("<B1-Motion>", self._on_drag_motion)
        self.canvas.bind("<ButtonRelease-1>", self._on_drag_end)
    
    def load_schema(self, schema_data: Dict):
        """Load schema data and visualize"""
        self.tables = {}
        self.relationships = []
        
        # Extract tables
        if "tables" in schema_data:
            for table_name, table_info in schema_data["tables"].items():
                self.tables[table_name] = {
                    "name": table_name,
                    "fields": table_info.get("fields", {}),
                    "position": None  # Will be calculated
                }
        
        # Extract relationships
        if "relationships" in schema_data:
            self.relationships = schema_data["relationships"]
        
        self.redraw()
    
    def load_schema_from_recommendations(self, recommendations: Dict):
        """Load schema from migration recommendations"""
        self.tables = {}
        self.relationships = []
        
        # Extract from recommendations structure
        if "tables" in recommendations:
            tables_data = recommendations["tables"]
            for table_name, table_def in tables_data.items():
                fields = {}
                if "fields" in table_def:
                    for field_name, field_type in table_def["fields"].items():
                        fields[field_name] = {
                            "type": field_type,
                            "primary_key": field_name.endswith("_id") or field_name == "id"
                        }
                
                self.tables[table_name] = {
                    "name": table_name,
                    "fields": fields,
                    "position": None
                }
        
        # Try to infer relationships from field names
        self._infer_relationships()
        
        self.redraw()
    
    def _infer_relationships(self):
        """Infer relationships from field names"""
        for table_name, table_data in self.tables.items():
            for field_name, field_info in table_data["fields"].items():
                # Look for foreign key patterns
                if field_name.endswith("_id") and field_name != "id":
                    potential_table = field_name[:-3] + "s"  # user_id -> users
                    if potential_table in self.tables:
                        self.relationships.append({
                            "from_table": table_name,
                            "from_field": field_name,
                            "to_table": potential_table,
                            "to_field": "id",
                            "type": "many-to-one",
                            "confidence": 80
                        })
    
    def redraw(self):
        """Redraw the entire visualization"""
        self.canvas.delete("all")
        
        if not self.tables:
            self.canvas.create_text(
                400, 300,
                text="No schema loaded\nRun migration to visualize schema",
                font=("Arial", 14),
                fill=self.theme.get_color("fg_secondary"),
                justify=tk.CENTER
            )
            return
        
        # Calculate positions
        self._calculate_positions()
        
        # Draw relationships first (so they appear behind tables)
        self._draw_relationships()
        
        # Draw tables
        for table_name, table_data in self.tables.items():
            self._draw_table(table_name, table_data)
        
        # Update scroll region
        self.canvas.config(scrollregion=self.canvas.bbox("all"))
    
    def _calculate_positions(self):
        """Calculate table positions based on layout"""
        layout = self.layout_var.get()
        num_tables = len(self.tables)
        
        if layout == "grid":
            cols = math.ceil(math.sqrt(num_tables))
            for i, (table_name, table_data) in enumerate(self.tables.items()):
                row = i // cols
                col = i % cols
                x = 100 + col * (self.table_width + self.table_padding)
                y = 100 + row * 200
                table_data["position"] = (x, y)
        
        elif layout == "circular":
            center_x = 400
            center_y = 300
            radius = 250
            angle_step = 2 * math.pi / num_tables
            
            for i, (table_name, table_data) in enumerate(self.tables.items()):
                angle = i * angle_step
                x = center_x + radius * math.cos(angle)
                y = center_y + radius * math.sin(angle)
                table_data["position"] = (x, y)
        
        else:  # auto
            # Simple left-to-right layout with relationship awareness
            x = 100
            y = 100
            for table_name, table_data in self.tables.items():
                table_data["position"] = (x, y)
                y += 200
                if y > 600:
                    y = 100
                    x += self.table_width + self.table_padding
    
    def _draw_table(self, table_name: str, table_data: Dict):
        """Draw a single table"""
        if not table_data["position"]:
            return
        
        x, y = table_data["position"]
        fields = table_data["fields"]
        
        # Calculate table height
        table_height = self.table_header_height + len(fields) * self.field_height
        
        # Draw table background
        self.canvas.create_rectangle(
            x, y, x + self.table_width, y + table_height,
            fill=self.theme.get_color("bg_secondary"),
            outline=self.theme.get_color("border"),
            width=2,
            tags=("table", table_name)
        )
        
        # Draw header
        self.canvas.create_rectangle(
            x, y, x + self.table_width, y + self.table_header_height,
            fill=self.theme.get_color("accent"),
            outline="",
            tags=("table_header", table_name)
        )
        
        self.canvas.create_text(
            x + self.table_width / 2, y + self.table_header_height / 2,
            text=table_name,
            font=("Arial", 11, "bold"),
            fill=self.theme.get_color("button_fg"),
            tags=("table_name", table_name)
        )
        
        # Draw fields
        field_y = y + self.table_header_height
        for field_name, field_info in fields.items():
            # Field background
            if isinstance(field_info, dict):
                is_pk = field_info.get("primary_key", False)
                field_type = field_info.get("type", "")
            else:
                is_pk = field_name == "id"
                field_type = str(field_info)
            
            # Highlight primary keys (use lighter shade without alpha)
            if is_pk:
                pk_color = "#E8F5E9" if self.theme.theme_name == "light" else "#1B5E20"
                self.canvas.create_rectangle(
                    x + 2, field_y + 2,
                    x + self.table_width - 2, field_y + self.field_height - 2,
                    fill=pk_color,
                    outline="",
                    tags=("field_bg", table_name)
                )
            
            # Field name
            self.canvas.create_text(
                x + 10, field_y + self.field_height / 2,
                text=f"{'🔑 ' if is_pk else ''}{field_name}",
                font=("Arial", 9, "bold" if is_pk else "normal"),
                fill=self.theme.get_color("fg"),
                anchor=tk.W,
                tags=("field", table_name)
            )
            
            # Field type
            if field_type:
                self.canvas.create_text(
                    x + self.table_width - 10, field_y + self.field_height / 2,
                    text=field_type[:15],
                    font=("Arial", 8),
                    fill=self.theme.get_color("fg_secondary"),
                    anchor=tk.E,
                    tags=("field_type", table_name)
                )
            
            field_y += self.field_height
    
    def _draw_relationships(self):
        """Draw relationship lines between tables"""
        for rel in self.relationships:
            from_table = rel.get("from_table")
            to_table = rel.get("to_table")
            
            if from_table not in self.tables or to_table not in self.tables:
                continue
            
            from_pos = self.tables[from_table].get("position")
            to_pos = self.tables[to_table].get("position")
            
            if not from_pos or not to_pos:
                continue
            
            # Calculate connection points
            from_x = from_pos[0] + self.table_width
            from_y = from_pos[1] + self.table_header_height / 2
            to_x = to_pos[0]
            to_y = to_pos[1] + self.table_header_height / 2
            
            # Draw curved line
            self._draw_connection_line(from_x, from_y, to_x, to_y, rel)
    
    def _draw_connection_line(self, x1, y1, x2, y2, rel):
        """Draw a curved connection line"""
        # Calculate control points for bezier curve
        dx = x2 - x1
        dy = y2 - y1
        
        # Control points for smooth curve
        cx1 = x1 + dx * 0.5
        cy1 = y1
        cx2 = x2 - dx * 0.5
        cy2 = y2
        
        # Draw line
        confidence = rel.get("confidence", 100)
        color = self.theme.get_color("accent") if confidence > 70 else self.theme.get_color("warning")
        
        self.canvas.create_line(
            x1, y1, cx1, cy1, cx2, cy2, x2, y2,
            smooth=True,
            fill=color,
            width=2,
            arrow=tk.LAST,
            tags=("relationship",)
        )
        
        # Add label at midpoint
        mid_x = (x1 + x2) / 2
        mid_y = (y1 + y2) / 2
        
        rel_type = rel.get("type", "")
        if rel_type:
            self.canvas.create_text(
                mid_x, mid_y - 10,
                text=rel_type,
                font=("Arial", 8),
                fill=self.theme.get_color("fg_secondary"),
                tags=("rel_label",)
            )
    
    def zoom_in(self):
        """Zoom in"""
        self.zoom_level *= 1.2
        self.canvas.scale("all", 400, 300, 1.2, 1.2)
    
    def zoom_out(self):
        """Zoom out"""
        self.zoom_level /= 1.2
        self.canvas.scale("all", 400, 300, 1/1.2, 1/1.2)
    
    def reset_view(self):
        """Reset zoom and pan"""
        self.zoom_level = 1.0
        self.pan_offset = [0, 0]
        self.redraw()
    
    def export_image(self):
        """Export visualization as PNG"""
        try:
            from PIL import Image, ImageDraw
            import tkinter as tk
            
            # Get canvas bbox
            x1, y1, x2, y2 = self.canvas.bbox("all")
            
            if not x1:
                messagebox.showwarning("Export", "No content to export")
                return
            
            # Ask for save location
            filename = filedialog.asksaveasfilename(
                defaultextension=".png",
                filetypes=[("PNG Image", "*.png"), ("All Files", "*.*")]
            )
            
            if not filename:
                return
            
            # Use canvas postscript (basic export)
            self.canvas.postscript(file=filename + ".ps")
            messagebox.showinfo("Export", f"Schema exported (as PostScript)\nNote: Install Pillow for PNG export")
            
        except ImportError:
            messagebox.showinfo("Export", 
                              "PNG export requires Pillow library\nInstall with: pip install Pillow")
        except Exception as e:
            messagebox.showerror("Export Error", str(e))
    
    def _on_mousewheel(self, event):
        """Handle mouse wheel for zooming"""
        if event.delta > 0:
            self.zoom_in()
        else:
            self.zoom_out()
    
    def _on_drag_start(self, event):
        """Start dragging"""
        self.canvas.scan_mark(event.x, event.y)
        self.dragging = True
    
    def _on_drag_motion(self, event):
        """Handle drag motion"""
        if self.dragging:
            self.canvas.scan_dragto(event.x, event.y, gain=1)
    
    def _on_drag_end(self, event):
        """End dragging"""
        self.dragging = False


if __name__ == "__main__":
    # Test visualization
    from theme_manager import ThemeManager
    
    root = tk.Tk()
    root.title("Schema Visualizer Test")
    root.geometry("900x700")
    
    theme = ThemeManager(root, "light")
    
    viz = SchemaVisualizer(root, theme)
    
    # Test data
    test_schema = {
        "tables": {
            "users": {
                "fields": {
                    "id": {"type": "INT", "primary_key": True},
                    "name": {"type": "VARCHAR(100)", "primary_key": False},
                    "email": {"type": "VARCHAR(255)", "primary_key": False},
                    "age": {"type": "INT", "primary_key": False}
                }
            },
            "posts": {
                "fields": {
                    "id": {"type": "INT", "primary_key": True},
                    "title": {"type": "VARCHAR(200)", "primary_key": False},
                    "user_id": {"type": "INT", "primary_key": False},
                    "content": {"type": "TEXT", "primary_key": False}
                }
            },
            "comments": {
                "fields": {
                    "id": {"type": "INT", "primary_key": True},
                    "post_id": {"type": "INT", "primary_key": False},
                    "user_id": {"type": "INT", "primary_key": False},
                    "text": {"type": "TEXT", "primary_key": False}
                }
            }
        },
        "relationships": [
            {
                "from_table": "posts",
                "from_field": "user_id",
                "to_table": "users",
                "to_field": "id",
                "type": "many-to-one",
                "confidence": 95
            },
            {
                "from_table": "comments",
                "from_field": "post_id",
                "to_table": "posts",
                "to_field": "id",
                "type": "many-to-one",
                "confidence": 95
            }
        ]
    }
    
    viz.load_schema(test_schema)
    
    root.mainloop()