import os
import sys
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch

# Add project root to sys.path
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(PROJECT_ROOT.parent))
sys.path.append(str(PROJECT_ROOT))

# Setup dummy bpy classes before importing NodeSlicer
from unittest.mock import MagicMock
class DummyNodeSocket: pass
class DummyNode: pass
class DummyObject: pass
class DummyMesh: pass
class DummyNodeSocketShader: pass
class DummyNodeCustomGroup: pass
class DummyNodeTree: pass
class DummyOperator: pass
class DummyMenu: pass
class DummyPanel: pass
class DummyPropertyGroup: pass

mock_types = MagicMock()
mock_types.NodeSocket = DummyNodeSocket
mock_types.Node = DummyNode
mock_types.Object = DummyObject
mock_types.Mesh = DummyMesh
mock_types.NodeSocketShader = DummyNodeSocketShader
mock_types.NodeCustomGroup = DummyNodeCustomGroup
mock_types.NodeTree = DummyNodeTree
mock_types.Operator = DummyOperator
mock_types.Menu = DummyMenu
mock_types.Panel = DummyPanel
mock_types.PropertyGroup = DummyPropertyGroup

mock_bpy = MagicMock()
mock_bpy.context.scene = MagicMock()
mock_bpy.context.scene.unit_settings.scale_length = 0.001
mock_bpy.types = mock_types

sys.modules['bpy'] = mock_bpy
sys.modules['bpy.types'] = mock_types
sys.modules['bpy.props'] = MagicMock()
sys.modules['bpy.utils'] = MagicMock()
sys.modules['bpy.utils.previews'] = MagicMock()
sys.modules['bpy.app'] = MagicMock()
sys.modules['bpy.app.handlers'] = MagicMock()

from NodeSlicer.core.slice import SlicePaths, is_closed_path

VISUAL_DIR = Path(__file__).resolve().parent
IMAGES_DIR = VISUAL_DIR / "images"
IMAGES_DIR.mkdir(parents=True, exist_ok=True)


def draw_path_with_arrows(ax, path, color, label=None, linewidth=2.0, linestyle='-', marker='o', markersize=4, alpha=1.0):
    xs = path[:, 0]
    ys = path[:, 1]
    line, = ax.plot(xs, ys, color=color, linewidth=linewidth, linestyle=linestyle, marker=marker, markersize=markersize, alpha=alpha, label=label)
    
    # Add small directional arrows along segments
    for i in range(len(path) - 1):
        x0, y0 = path[i, 0], path[i, 1]
        x1, y1 = path[i + 1, 0], path[i + 1, 1]
        dx = x1 - x0
        dy = y1 - y0
        seg_len = np.hypot(dx, dy)
        if seg_len > 0.5:
            mid_x = (x0 + x1) / 2.0
            mid_y = (y0 + y1) / 2.0
            arrow_dx = (dx / seg_len) * 0.4
            arrow_dy = (dy / seg_len) * 0.4
            ax.annotate(
                "",
                xy=(mid_x + arrow_dx, mid_y + arrow_dy),
                xytext=(mid_x - arrow_dx, mid_y - arrow_dy),
                arrowprops=dict(arrowstyle="->", color=color, lw=1.5, mutation_scale=12)
            )
    return line


def generate_single_line_comparison() -> str:
    fig, axes = plt.subplots(1, 3, figsize=(16, 5.5), dpi=150)
    
    # 10x10 square boundary
    square = np.array([
        [0.0, 0.0, 0.2],
        [10.0, 0.0, 0.2],
        [10.0, 10.0, 0.2],
        [0.0, 10.0, 0.2],
        [0.0, 0.0, 0.2],
    ], dtype=np.float64)
    boundary_sp = SlicePaths([square])
    
    infill = np.array([[5.0, 0.0, 0.2], [5.0, 10.0, 0.2]], dtype=np.float64)
    infill_sp = SlicePaths([infill])
    
    # Subplot 1: Inputs
    ax1 = axes[0]
    ax1.set_title("1. Input Geometry\n(Open Infill Line + Closed Perimeter)", fontsize=11, fontweight='bold', pad=10)
    draw_path_with_arrows(ax1, square, color='#1f77b4', label="Anchor Boundary (Input 2)", linewidth=2.0)
    draw_path_with_arrows(ax1, infill, color='#ff7f0e', label="Open Path (Input 1)", linewidth=2.5)
    ax1.plot(infill[0, 0], infill[0, 1], 'go', markersize=9, label="Start Vertex P[0]")
    ax1.plot(infill[1, 0], infill[1, 1], 'ro', markersize=9, label="End Vertex P[-1]")
    ax1.set_xlim(-1, 11)
    ax1.set_ylim(-1, 11)
    ax1.set_aspect('equal')
    ax1.grid(True, linestyle=':', alpha=0.6)
    ax1.legend(loc='lower center', fontsize=8)
    
    # Subplot 2: Reverse = False
    res_fwd = infill_sp.anchor(boundary_sp, distance=2.5, reverse=False)
    p_fwd = res_fwd[0]
    ax2 = axes[1]
    ax2.set_title("2. Anchored (Reverse = False)\nFollows Forward (+1) Traversal", fontsize=11, fontweight='bold', pad=10)
    ax2.plot(square[:, 0], square[:, 1], color='#1f77b4', linestyle='--', linewidth=1.5, alpha=0.5, label="Anchor Perimeter")
    # Draw original in light grey
    ax2.plot(infill[:, 0], infill[:, 1], color='gray', linestyle=':', linewidth=1.5, alpha=0.6)
    # Draw anchored path with directional arrows
    draw_path_with_arrows(ax2, p_fwd, color='#2ca02c', label="Anchored Path", linewidth=2.5)
    # Highlight start anchor and end anchor
    ax2.plot(p_fwd[:2, 0], p_fwd[:2, 1], color='#9467bd', linewidth=3.5, label="Start Anchor (Prepended)")
    ax2.plot(p_fwd[-2:, 0], p_fwd[-2:, 1], color='#d62728', linewidth=3.5, label="End Anchor (Appended)")
    ax2.plot(p_fwd[0, 0], p_fwd[0, 1], 'g^', markersize=10, label="New Start (2.5mm before)")
    ax2.plot(p_fwd[-1, 0], p_fwd[-1, 1], 'rs', markersize=9, label="New End (2.5mm after)")
    ax2.set_xlim(-1, 11)
    ax2.set_ylim(-1, 11)
    ax2.set_aspect('equal')
    ax2.grid(True, linestyle=':', alpha=0.6)
    ax2.legend(loc='lower center', fontsize=8)
    
    # Subplot 3: Reverse = True
    res_rev = infill_sp.anchor(boundary_sp, distance=2.5, reverse=True)
    p_rev = res_rev[0]
    ax3 = axes[2]
    ax3.set_title("3. Anchored (Reverse = True)\nFollows Opposite (-1) Traversal", fontsize=11, fontweight='bold', pad=10)
    ax3.plot(square[:, 0], square[:, 1], color='#1f77b4', linestyle='--', linewidth=1.5, alpha=0.5, label="Anchor Perimeter")
    ax3.plot(infill[:, 0], infill[:, 1], color='gray', linestyle=':', linewidth=1.5, alpha=0.6)
    draw_path_with_arrows(ax3, p_rev, color='#e377c2', label="Anchored Path", linewidth=2.5)
    ax3.plot(p_rev[:2, 0], p_rev[:2, 1], color='#9467bd', linewidth=3.5, label="Start Anchor (Prepended)")
    ax3.plot(p_rev[-2:, 0], p_rev[-2:, 1], color='#d62728', linewidth=3.5, label="End Anchor (Appended)")
    ax3.plot(p_rev[0, 0], p_rev[0, 1], 'g^', markersize=10, label="New Start (2.5mm before)")
    ax3.plot(p_rev[-1, 0], p_rev[-1, 1], 'rs', markersize=9, label="New End (2.5mm after)")
    ax3.set_xlim(-1, 11)
    ax3.set_ylim(-1, 11)
    ax3.set_aspect('equal')
    ax3.grid(True, linestyle=':', alpha=0.6)
    ax3.legend(loc='lower center', fontsize=8)
    
    plt.tight_layout()
    out_path = IMAGES_DIR / "anchor_paths_comparison.png"
    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    return str(out_path)


def generate_infill_hatch_visualization() -> str:
    fig, axes = plt.subplots(1, 2, figsize=(14, 6.5), dpi=150)
    
    # Circular perimeter (radius 10)
    angles = np.linspace(0, 2 * np.pi, 65)
    circle_pts = np.column_stack([10 * np.cos(angles), 10 * np.sin(angles), np.full(len(angles), 0.2)])
    boundary_sp = SlicePaths([circle_pts])
    
    # Generate parallel horizontal infill lines clipped to circle
    lines = []
    for y in np.arange(-8.0, 9.0, 2.0):
        half_w = np.sqrt(max(0.0, 100.0 - y * y))
        lines.append(np.array([[-half_w, y, 0.2], [half_w, y, 0.2]], dtype=np.float64))
    
    infill_sp = SlicePaths(lines)
    
    # Subplot 1: Without anchor
    ax1 = axes[0]
    ax1.set_title("Standard Infill Hatch\nEndpoints touching circular perimeter", fontsize=11, fontweight='bold', pad=10)
    ax1.plot(circle_pts[:, 0], circle_pts[:, 1], color='#1f77b4', linewidth=2.5, label="Boundary Perimeter")
    for i, line in enumerate(lines):
        ax1.plot(line[:, 0], line[:, 1], color='#ff7f0e', linewidth=1.8, label="Infill Lines" if i == 0 else None)
        ax1.plot(line[0, 0], line[0, 1], 'g.', markersize=7)
        ax1.plot(line[-1, 0], line[-1, 1], 'r.', markersize=7)
    ax1.set_xlim(-12, 12)
    ax1.set_ylim(-12, 12)
    ax1.set_aspect('equal')
    ax1.grid(True, linestyle=':', alpha=0.6)
    ax1.legend(loc='upper right', fontsize=9)
    
    # Subplot 2: With anchor (distance = 1.5mm)
    res_anchored = infill_sp.anchor(boundary_sp, distance=1.5, reverse=False)
    ax2 = axes[1]
    ax2.set_title("Anchored Infill Hatch (Distance = 1.5mm)\nEndpoints follow circular perimeter", fontsize=11, fontweight='bold', pad=10)
    ax2.plot(circle_pts[:, 0], circle_pts[:, 1], color='#1f77b4', linestyle='--', linewidth=1.5, alpha=0.5, label="Boundary Perimeter")
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22']
    for i, p in enumerate(res_anchored):
        c = colors[i % len(colors)]
        draw_path_with_arrows(ax2, p, color=c, linewidth=2.0, markersize=3, label="Anchored Pass" if i == 0 else None)
    ax2.set_xlim(-12, 12)
    ax2.set_ylim(-12, 12)
    ax2.set_aspect('equal')
    ax2.grid(True, linestyle=':', alpha=0.6)
    ax2.legend(loc='upper right', fontsize=9)
    
    plt.tight_layout()
    out_path = IMAGES_DIR / "anchor_paths_infill_hatch.png"
    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    return str(out_path)


def generate_corner_wrapping_visualization() -> str:
    fig, ax = plt.subplots(figsize=(8, 7), dpi=150)
    
    square = np.array([
        [0.0, 0.0, 0.2],
        [10.0, 0.0, 0.2],
        [10.0, 10.0, 0.2],
        [0.0, 10.0, 0.2],
        [0.0, 0.0, 0.2],
    ], dtype=np.float64)
    boundary_sp = SlicePaths([square])
    
    # Infill starts at (1.0, 0.0) - only 1mm from corner (0,0) - and goes to (1.0, 10.0)
    infill = np.array([[1.0, 0.0, 0.2], [1.0, 10.0, 0.2]], dtype=np.float64)
    infill_sp = SlicePaths([infill])
    
    res = infill_sp.anchor(boundary_sp, distance=3.5, reverse=False)
    p = res[0]
    
    ax.set_title("Anchor Paths: Corner Turning Verification\n(Distance = 3.5mm wraps around (0,0) corner)", fontsize=11, fontweight='bold', pad=10)
    ax.plot(square[:, 0], square[:, 1], color='#1f77b4', linestyle='--', linewidth=2.0, label="Square Boundary Perimeter")
    ax.plot(infill[:, 0], infill[:, 1], color='gray', linestyle=':', linewidth=2.0, label="Original Infill Line")
    
    draw_path_with_arrows(ax, p, color='#2ca02c', linewidth=2.5, markersize=6, label="Anchored Toolpath")
    
    # Highlight corner points
    ax.plot(p[0, 0], p[0, 1], 'g^', markersize=12, label="Start Point on Left Edge (0, 2.5)")
    ax.plot(p[1, 0], p[1, 1], 'mo', markersize=10, label="Corner Point Traversed (0, 0)")
    ax.plot(p[2, 0], p[2, 1], 'ys', markersize=9, label="Original Start Point (1, 0)")
    
    # Annotate distance
    ax.annotate("Corner turned!\n2.5mm up left edge\n+ 1.0mm along bottom = 3.5mm", xy=(0.0, 1.25), xytext=(2.0, 2.5),
                arrowprops=dict(facecolor='black', shrink=0.05, width=1, headwidth=6),
                fontsize=9, fontweight='semibold')
    
    ax.set_xlim(-1.5, 11)
    ax.set_ylim(-1.5, 11)
    ax.set_aspect('equal')
    ax.grid(True, linestyle=':', alpha=0.6)
    ax.legend(loc='upper right', fontsize=8.5)
    
    plt.tight_layout()
    out_path = IMAGES_DIR / "anchor_paths_corners.png"
    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    return str(out_path)


def generate_partial_touch_visualization() -> str:
    fig, axes = plt.subplots(1, 3, figsize=(16, 5), dpi=150)
    
    square = np.array([
        [0.0, 0.0, 0.2],
        [10.0, 0.0, 0.2],
        [10.0, 10.0, 0.2],
        [0.0, 10.0, 0.2],
        [0.0, 0.0, 0.2],
    ], dtype=np.float64)
    boundary_sp = SlicePaths([square])
    
    # Case A: Start-only touch
    infill_start = np.array([[5.0, 0.0, 0.2], [5.0, 5.0, 0.2]], dtype=np.float64)
    sp_a = SlicePaths([infill_start])
    res_a = sp_a.anchor(boundary_sp, distance=2.0)
    
    ax1 = axes[0]
    ax1.set_title("Start-Only Touch\n(Start touches bottom, end in interior)", fontsize=10.5, fontweight='bold')
    ax1.plot(square[:, 0], square[:, 1], 'b--', alpha=0.5, label="Boundary")
    draw_path_with_arrows(ax1, res_a[0], color='#2ca02c', label="Anchored Path")
    ax1.plot(res_a[0][0, 0], res_a[0][0, 1], 'g^', markersize=9, label="Prepended Start")
    ax1.plot(infill_start[-1, 0], infill_start[-1, 1], 'rx', markersize=9, label="Untouched End")
    ax1.set_xlim(-1, 11)
    ax1.set_ylim(-1, 11)
    ax1.set_aspect('equal')
    ax1.grid(True, linestyle=':', alpha=0.6)
    ax1.legend(loc='lower center', fontsize=8)
    
    # Case B: End-only touch
    infill_end = np.array([[5.0, 5.0, 0.2], [5.0, 10.0, 0.2]], dtype=np.float64)
    sp_b = SlicePaths([infill_end])
    res_b = sp_b.anchor(boundary_sp, distance=2.0)
    
    ax2 = axes[1]
    ax2.set_title("End-Only Touch\n(Start in interior, end touches top)", fontsize=10.5, fontweight='bold')
    ax2.plot(square[:, 0], square[:, 1], 'b--', alpha=0.5, label="Boundary")
    draw_path_with_arrows(ax2, res_b[0], color='#ff7f0e', label="Anchored Path")
    ax2.plot(infill_end[0, 0], infill_end[0, 1], 'gx', markersize=9, label="Untouched Start")
    ax2.plot(res_b[0][-1, 0], res_b[0][-1, 1], 'rs', markersize=9, label="Appended End")
    ax2.set_xlim(-1, 11)
    ax2.set_ylim(-1, 11)
    ax2.set_aspect('equal')
    ax2.grid(True, linestyle=':', alpha=0.6)
    ax2.legend(loc='lower center', fontsize=8)
    
    # Case C: Closed path (ignored)
    box = np.array([[3.0, 3.0, 0.2], [7.0, 3.0, 0.2], [7.0, 7.0, 0.2], [3.0, 7.0, 0.2], [3.0, 3.0, 0.2]], dtype=np.float64)
    sp_c = SlicePaths([box])
    res_c = sp_c.anchor(boundary_sp, distance=2.0)
    
    ax3 = axes[2]
    ax3.set_title("Closed Path Ignored\n(Closed loops remain untouched)", fontsize=10.5, fontweight='bold')
    ax3.plot(square[:, 0], square[:, 1], 'b--', alpha=0.5, label="Boundary")
    draw_path_with_arrows(ax3, res_c[0], color='#9467bd', label="Untouched Closed Loop")
    ax3.set_xlim(-1, 11)
    ax3.set_ylim(-1, 11)
    ax3.set_aspect('equal')
    ax3.grid(True, linestyle=':', alpha=0.6)
    ax3.legend(loc='lower center', fontsize=8)
    
    plt.tight_layout()
    out_path = IMAGES_DIR / "anchor_paths_partial_touch.png"
    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    return str(out_path)


def run_all_anchor_visualizations() -> list[str]:
    print("Generating visual examples for Anchor Paths...")
    p1 = generate_single_line_comparison()
    print("Generated:", p1)
    p2 = generate_infill_hatch_visualization()
    print("Generated:", p2)
    p3 = generate_corner_wrapping_visualization()
    print("Generated:", p3)
    p4 = generate_partial_touch_visualization()
    print("Generated:", p4)
    return [p1, p2, p3, p4]


if __name__ == '__main__':
    run_all_anchor_visualizations()