import json
import sys
import math

from kikit.common import findBoardBoundingBox
from kikit.panelize_ui_impl import buildText
from kikit.panelize_ui_sections import ppText
from kikit.defs import *
from kikit.units import mm
from pcbnewTransition import pcbnew
from pcbnewTransition.pcbnew import LoadBoard

PCB_DIV_MM = 1000000

pcb_file = sys.argv[1]
panel_file = sys.argv[2]

# Origin mark settings
ORIGIN_SQUARE_SIZE = 2.0  # mm
ORIGIN_CIRCLE_DIAMETER = 3.0  # mm
ORIGIN_CIRCLE_WIDTH = 0.4  # mm line width

def str_mm_to_float(input: str) -> float:
    input = input.lower()
    # Panic, should all be in mm for sanity's sake
    if "mm" not in input:
        print(input, "not in mm")
        exit(1)
    return float(input.replace("mm", ""))

def float_to_str_mm(input: float, n: int=0) -> str:
    return "{:.2f}".format(input).zfill(n+3)+"mm"

def float_to_deg(input: float, n: int=0) -> str:
    return "{:.2f}".format(input).zfill(n+3)+"deg"


class Text:
    def __init__(self, text: str, anchor: str, offset: tuple[float, float], rot: float, size: tuple[float, float, float] = (1.0, 1.0, 0.16), just: tuple[str, str] = ('center', 'center')):
        self.layer=Layer.F_SilkS
        self.text = text
        self.anchor = anchor
        self.hoff = float_to_str_mm(offset[0])
        self.voff = float_to_str_mm(offset[1])
        self.rot = str(rot)+"deg"
        self.width = float_to_str_mm(size[0])
        self.height = float_to_str_mm(size[1])
        self.thickness = float_to_str_mm(size[2])
        self.hjust = just[0]
        self.vjust = just[1]

    def flip(self):
        self.layer=Layer.B_SilkS
        rot = float(self.rot.replace("deg", ""))
        rot += 180.0 
        self.rot = str(rot)+"deg"
        return self

    def to_dict(self, n: int) -> dict:
        if n == 0:
            str_n = ''
        else:
            str_n = str(n + 1)
        return {
            "text" + str_n: 
            {
                "type": "simple",
                "text": self.text,
                "anchor": self.anchor,
                "hoffset": self.hoff,
                "voffset": self.voff,
                "orientation": self.rot,
                "width": self.width,
                "height": self.height,
                "hjustify": self.hjust,
                "vjustify": self.vjust,
                "thickness": self.thickness,
                "layer": int(self.layer),
            }
        }


class Frame:
    def __init__(self, ftype: str, width: str, spacing: str):
        self.ftype = ftype.lower()
        self.width = str_mm_to_float(width)
        self.spacing = str_mm_to_float(spacing)
        if "frame" in self.ftype:
            self.extra_d = [(self.width + self.spacing)*2, (self.width + self.spacing)*2]
        elif "tb" in self.ftype:
            self.extra_d = [0, (self.width + self.spacing)*2]
        else:
            self.extra_d = [(self.width + self.spacing)*2, 0]

class Board:
    def __init__(self, fids: list[tuple[float, float]], rot: float, dims: tuple[float, float]):
        fid_list = []
        for f in fids:
            fid_list.append("(" + float_to_str_mm(f[0]) + ", " + float_to_str_mm(f[1]) + ")")
        self.fids = ",".join(fid_list)
        self.rot = rot
        self.dims = dims


class Panel:
    def __init__(self, spacing: str, grid: tuple[int, int], frame: Frame, board: Board):
        self.spacing = str_mm_to_float(spacing)
        self.grid = grid
        self.gaps = (grid[0] - 1, grid[1] - 1)
        self.frame = frame
        self.text = {}
        self.dims = [self.gaps[i]*self.spacing + self.grid[i]*board.dims[i] + frame.extra_d[i] for i in range(0,2)]
        self.spacing = [self.spacing + board.dims[i] for i in range(0,2)]

        arrow_off = [frame.spacing + frame.width, frame.width]
        text_off = [arrow_off[0] + 2, frame.width/2]

        if "frame" in frame.ftype or "tb" in frame.ftype:
            txt_loc = ('mt', 'bl') 
            txt_off = ((0, text_off[1]), (text_off[0], -text_off[1]))
            txt_just = (('center', 'center'), ('left', 'center'))
            txt_rot = 0 
        else:
            txt_loc = ('ml', 'br') 
            txt_off = ((text_off[1], 0), (-text_off[1], -text_off[0]))
            txt_just = (('center', 'center'), ('left', 'center'))
            txt_rot = 90

        self.text.update(Text("{boardTitle}-{boardRevision}", txt_loc[0], txt_off[0], txt_rot, just=txt_just[0]).to_dict(len(self.text)))

        spacer = "     "
        txt =\
        "PNL (x,y):   (" + float_to_str_mm(self.dims[0],4) + ", " + float_to_str_mm(self.dims[1],4) + ")\n" +\
        "SPA (x,y):   (" + float_to_str_mm(self.spacing[0],2) + ", " + float_to_str_mm(self.spacing[1],2) +")\n" +\
        "BRD (x,y),r: (" + float_to_str_mm(board.dims[0],4) + ", " + float_to_str_mm(board.dims[1],4) + "), " +\
            str(board.rot) + "\n" +\
        "FIDS (x,y) : " + board.fids
        self.text.update(Text(txt, txt_loc[1], txt_off[1], txt_rot, just=txt_just[1]).to_dict(1))

        if "frame" in self.frame.ftype:
            txt_orig = [
Text("ORIGX ->", "bl", (arrow_off[0], -arrow_off[1]), 90, just=("right", "center")),
Text("ORIGY ->", "bl", (arrow_off[1], -arrow_off[0]), 0, just=("right", "center"))
            ]
            for txt in txt_orig:
                self.text.update(txt.to_dict(len(self.text)))
                self.text.update(txt.flip().to_dict(len(self.text)))

def addOriginMark(board, position, rotation=0):
    """
    Add origin mark at the given position:
    - 2x2mm filled copper region on F.Cu and B.Cu with corner at origin
    - 4mm diameter silkscreen circle centered on origin

    position: pcbnew.VECTOR2I in internal units
    rotation: rotation in degrees
    """
    square_size = int(ORIGIN_SQUARE_SIZE * PCB_DIV_MM)
    circle_radius = int((ORIGIN_CIRCLE_DIAMETER / 2) * PCB_DIV_MM)
    line_width = int(ORIGIN_CIRCLE_WIDTH * PCB_DIV_MM)

    # Create filled copper polygon for each copper layer
    for layer in [pcbnew.F_Cu, pcbnew.B_Cu]:
        poly = pcbnew.PCB_SHAPE(board)
        poly.SetShape(pcbnew.SHAPE_T_POLY)
        poly.SetLayer(layer)
        poly.SetFilled(True)
        poly.SetWidth(0)  # No outline, just fill

        # Create square polygon with corner at origin
        # Square extends in -X, -Y direction (away from board, into frame area)
        # For B.Cu, mirror in X direction
        if layer == pcbnew.B_Cu:
            corners = [
                pcbnew.VECTOR2I(position.x, position.y),
                pcbnew.VECTOR2I(position.x + square_size, position.y),
                pcbnew.VECTOR2I(position.x + square_size, position.y - square_size),
                pcbnew.VECTOR2I(position.x, position.y - square_size),
            ]
        else:
            corners = [
                pcbnew.VECTOR2I(position.x, position.y),
                pcbnew.VECTOR2I(position.x - square_size, position.y),
                pcbnew.VECTOR2I(position.x - square_size, position.y - square_size),
                pcbnew.VECTOR2I(position.x, position.y - square_size),
            ]

        # Set the polygon points
        poly.SetPolyPoints(corners)
        board.Add(poly)

    # Add silkscreen circles on F.SilkS and B.SilkS
    for layer in [pcbnew.F_SilkS, pcbnew.B_SilkS]:
        circle = pcbnew.PCB_SHAPE(board)
        circle.SetShape(pcbnew.SHAPE_T_CIRCLE)
        circle.SetLayer(layer)
        circle.SetCenter(position)
        circle.SetEnd(pcbnew.VECTOR2I(position.x + circle_radius, position.y))
        circle.SetWidth(line_width)
        board.Add(circle)

    # Add circular keepout zone to prevent copper pour (except our rectangle)
    # Create a circular polygon approximation for the keepout
    num_segments = 32
    for layer in [pcbnew.F_Cu, pcbnew.B_Cu]:
        keepout = pcbnew.ZONE(board)
        keepout.SetIsRuleArea(True)
        keepout.SetDoNotAllowCopperPour(True)
        keepout.SetDoNotAllowTracks(False)
        keepout.SetDoNotAllowVias(False)
        keepout.SetDoNotAllowPads(False)
        keepout.SetDoNotAllowFootprints(False)
        keepout.SetLayer(layer)

        # Create circular outline
        outline = keepout.Outline()
        outline.NewOutline()
        for i in range(num_segments):
            angle = 2 * math.pi * i / num_segments
            x = position.x + int(circle_radius * math.cos(angle))
            y = position.y + int(circle_radius * math.sin(angle))
            outline.Append(x, y)

        board.Add(keepout)


def isMarkSafe(mark_center, mark_radius, substrates):
    """
    Check if a mark center is at a valid position (on edge/corner or outside boards).
    Marks at board corners are allowed - part of the circle may overlap the board edge.
    Only reject if the mark center is fully INSIDE a board (not at an edge).
    """
    EDGE_TOLERANCE = int(0.5 * PCB_DIV_MM)  # 0.5mm tolerance for edge detection

    for substrate in substrates:
        bounds = substrate.bounds()  # (minx, miny, maxx, maxy) in internal units

        # Check if mark center is INSIDE the board (not on edge)
        # A point is "inside" if it's more than EDGE_TOLERANCE away from all edges
        inside_x = (bounds[0] + EDGE_TOLERANCE) < mark_center.x < (bounds[2] - EDGE_TOLERANCE)
        inside_y = (bounds[1] + EDGE_TOLERANCE) < mark_center.y < (bounds[3] - EDGE_TOLERANCE)

        if inside_x and inside_y:
            return False  # Mark center is inside this substrate - not safe

    return True  # Safe - mark is at edge/corner or outside all substrates


def kikitPostprocess(panel, arg):
    """
    KiKit post-processing callback.
    Adds origin marks at each board instance's origin in the panel.
    """
    print("Running post-panel processing...")

    # Parse the argument (contains origin offset from pre-processing)
    try:
        arg_data = json.loads(arg)
        origin_offset = arg_data.get("origin_offset", [0, 0])
    except (json.JSONDecodeError, TypeError):
        print("Warning: Could not parse script argument")
        origin_offset = [0, 0]

    print(f"Origin offset from bbox: ({origin_offset[0]:.2f}, {origin_offset[1]:.2f}) mm")

    # Get the panel board
    board = panel.board

    # Calculate mark radius for safety check (use circle radius + small margin)
    mark_check_radius = int((ORIGIN_CIRCLE_DIAMETER / 2 + 0.5) * PCB_DIV_MM)

    # Get all board substrates (instances) in the panel
    for i, substrate in enumerate(panel.substrates):
        # Use boundingBox() to get substrate position
        bbox = substrate.boundingBox()
        bbox_pos = bbox.GetPosition()

        # Calculate actual origin position (bbox corner + offset)
        origin_x = bbox_pos.x + int(origin_offset[0] * PCB_DIV_MM)
        origin_y = bbox_pos.y + int(origin_offset[1] * PCB_DIV_MM)
        origin = pcbnew.VECTOR2I(origin_x, origin_y)

        rotation = 0  # Substrates don't expose rotation directly

        print(f"Board {i}: actual origin at ({origin.x/PCB_DIV_MM:.2f}, {origin.y/PCB_DIV_MM:.2f}) mm")

        # Check if mark is safe to place (fully outside all boards)
        if isMarkSafe(origin, mark_check_radius, panel.substrates):
            print(f"  -> Safe to place, adding origin mark")
            addOriginMark(board, origin, rotation)
        else:
            print(f"  -> SKIPPING: mark would overlap with board area")

    print("Finished post panel script")

if __name__ == '__main__':
    print("Starting pre_panel script")

    board = LoadBoard(pcb_file)
    sourceArea = findBoardBoundingBox(board)
    fids = []
    for foot in board.Footprints():
        if "FID" in foot.GetReference():
            cent = foot.GetBoundingBox().GetCenter()
            fids.append(tuple([cent[i]/PCB_DIV_MM for i in range(0,2)]))

    # Get the board's auxiliary origin (design origin)
    aux_origin = board.GetDesignSettings().GetAuxOrigin()
    bbox_origin = pcbnew.VECTOR2I(sourceArea.GetX(), sourceArea.GetY())

    # Calculate offset from bounding box corner to aux origin (in mm)
    origin_offset_x = (aux_origin.x - bbox_origin.x) / PCB_DIV_MM
    origin_offset_y = (aux_origin.y - bbox_origin.y) / PCB_DIV_MM
    print(f"Board aux origin offset from bbox: ({origin_offset_x:.2f}, {origin_offset_y:.2f}) mm")

    json_file = open(panel_file)
    json_str = json_file.read()
    panel_json = json.loads(json_str)

    # Get spacing - support both "space" and "hspace"/"vspace" formats
    layout = panel_json["layout"]
    layout_space = layout.get("space", layout.get("hspace", "0mm"))

    # Get framing space with default
    framing = panel_json.get("framing", {})
    framing_space = framing.get("space", "0mm")
    framing_width = framing.get("width", "0mm")
    framing_type = framing.get("type", "none")

    frame = Frame(framing_type, framing_width, framing_space)
    board_obj = Board(fids, layout.get("rotation", "0deg"), (sourceArea.GetWidth()/PCB_DIV_MM, sourceArea.GetHeight()/PCB_DIV_MM))
    panel = Panel(layout_space, (int(layout["rows"]), int(layout["cols"])), frame, board_obj)

    import pathlib

    # Ensure post section exists
    if "post" not in panel_json:
        panel_json["post"] = {}

    panel_json["post"]["script"] = str(pathlib.Path(__file__).resolve())

    # Pass origin offset to post-processing for origin mark placement
    panel_json["post"]["scriptarg"] = json.dumps({
        "text": panel.text,
        "origin_offset": [origin_offset_x, origin_offset_y]
    })

    json_file = open(panel_file, mode="w")
    json.dump(panel_json, json_file, indent=4)

    print("Finished pre_panel script")