#!/bin/python
import math
import time
import os
import sys
import argparse
import logging

from pycbnew.primitives.project import KProject
from pycbnew.primitives.footprint import KFootprint, KFootprintProperty
from pycbnew.primitives.net import KNet
from pycbnew.utils.geometry import Point, Vector
from pycbnew.composites.gr_line_path import KGrLinePath
from pycbnew.primitives.gr_elements import (
    KGrRect,
    KGrPoly,
    Stroke,
    KGrText,
    Font,
    KGrDimension,
    DimensionFormat,
    DimensionStyle,
    TextJustify,
    KGrLine,
    KTable,
)
from pycbnew.primitives.segment import KSegment
import xml.etree.ElementTree as ET

# Global Variables set through argparse
LOG_LEVEL = logging.INFO
NO_POWER_NETS = False
NO_TEXT = False
NO_TRACES = False
SCALE = 1.0
PCB_FILE = "trainmap.kicad_pcb"
INPUT_XML = ""
MC_POS = Point(0, 0)
MARGIN = {"left": 10, "top": 5}
TRACE_WIDTH_MM=0.5
CAP_DISTANCE_MM=5
TEXT_SPACING_MM=15

# Fonts
TITLE_FONTFACE = ""
FONTFACE = ""
LINE_FONT = Font(size=(12, 12))
LINE_FONT.face = FONTFACE
TITLE_FONT = Font(size=(28, 28))
TITLE_FONT.face = TITLE_FONTFACE

# Footprints
KFootprint.set_footprint_folders([f"{os.path.dirname(__file__)}/kicad_libs"])
FOOTPRINT_C = "Capacitor_SMD:C_0603_1608Metric"
FOOTPRINT_LED = "Dialight_587_1032_147F:587"
FOOTPRINT_MC = "DFR0654:MODULE_DFR0654"

MC_PADS = {
    "GND": 30,
    "GND1": 4,
    "VCC": 32,
    "gpios": [4, 12, 13, 14, 15, 18, 19, 21, 22, 23, 25, 26, 34, 35, 2],
}
LED_PADS = {"GND": 0, "VCC": 1, "DI": 2, "DO": 3}

# Special Nets
NET_GND = KNet(1, "GND")
NET_VCC = KNet(2, "5V")

class Map:
    title: str
    lines: list
    width: int
    height: int

    def __init__(self, title):
        self.title = title
        self.lines = []
        self.width = 0
        self.height = 0

    def add(self, line):
        self.lines.append(line)


class Line:
    name: str
    color: str
    stations: list

    def __init__(self, name, color):
        self.name = name
        self.color = color
        self.stations = []

    def __str__(self):
        str = self.name
        if self.color != None:
            str += f" (color: {self.color})"
        return str

    def add(self, station):
        self.stations.append(station)


class Station:
    id: int
    name: str
    pos: Point
    angle: int
    text_pos: Point

    def __init__(self, id, name, x, y, angle=0, text_pos="right"):
        self.id = id
        self.name = name
        self.pos = Point(x * SCALE, y * SCALE)
        self.text_pos = text_pos
        self.angle = angle

    def __str__(self):
        return f"station: {self.name}, id: {self.id}, pos: ({str(self.pos.x)},{str(self.pos.y)}, angle: {self.angle} text_pos: {self.text_pos})"


"""
parses the xml with the map

syntax:
<map>
    <line name="%line_name%">
        <station x="%x_coord" y="%y_coord">%station_name</station>
    </line>
</map>
"""
def parse_map(xml_file):
    root = ET.parse(xml_file)

    title_node = root.find("title")
    title=""
    if title_node != None:
        title=title_node.text

    map = Map(title)
    logging.debug(f"map title: {title}")

    width = 0
    height = 0
    id=0
    for line in root.findall("line"):
        line_name = line.get("name")
        try:
            line_color = line.get("color")
        except Error:
            line_color = None

        line_parsed = Line(line_name, line_color)
        logging.debug(f"Parsing line: {line_parsed}")

        for station in line.findall("station"):
            id=id+1
            name = station.text
            x = int(float(station.get("x")))
            # search width
            if x > width:
                width = x
            y = int(float(station.get("y")))
            # search height
            if y > height:
                height = y

            station_parsed = Station(id, name, x, y)
            logging.debug(station_parsed)

            line_parsed.add(station_parsed)

        map.add(line_parsed)

    map.width = width
    map.height = height

    logging.debug(f"done parsing xml. pcb width: {
                  map.width}, height: {map.height}")

    return map

"""
Creates a capacity between LED's VCC and GND
"""
def create_capacity(station):
    cap_pos_x = CAP_DISTANCE_MM * math.sin(station.angle)
    cap_pos_y = CAP_DISTANCE_MM * math.cos(station.angle)
    cap_pos = Vector(cap_pos_x, cap_pos_y)
    fp = KFootprint(FOOTPRINT_C, at=station.pos + cap_pos, angle=station.angle)

    ref = fp.properties[KFootprintProperty.Type.Reference]
    ref.hide = 'yes'
    if not NO_POWER_NETS:
        fp.pads[0].net = NET_GND
        fp.pads[1].net = NET_VCC

    return fp

def calculate_angle(station1, station2):
    vec = Vector(station2.pos.x, station2.pos.y) - Vector(station1.pos.x, station1.pos.y) 
    angle_rad = vec.self_angle()
    return math.degrees(angle_rad)

"""
Creates the footprint for a station led
adds led, text and C
"""
def create_station_footprint(station):
    fp = KFootprint(FOOTPRINT_LED, at=station.pos, angle=station.angle)

    ref = fp.properties[KFootprintProperty.Type.Reference]
    if NO_TEXT:
        ref.layer = 'User.Comments'

    ref.value = station.name
    text_pos_x = TEXT_SPACING_MM * math.cos(station.angle)
    text_pos_y = TEXT_SPACING_MM * math.sin(station.angle)
    ref.at = Vector(text_pos_x, text_pos_y)

    if not NO_POWER_NETS:
        fp.pads[LED_PADS["GND"]].net = NET_GND
        fp.pads[LED_PADS["VCC"]].net = NET_VCC

    return fp


"""
Creates a new net between station 1 out pad and station 2 in pad
Connects them on the PCB
"""
def connect_stations(board, station1, station2, fp1, fp2):
    net = KNet(station1.id + 2, f"{station1.id}-{station2.id}")
    board.add(net)

    do = fp1.pads[LED_PADS["DO"]]
    di = fp2.pads[LED_PADS["DI"]]

    # connect the net
    do.net = net
    di.net = net

    if NO_TRACES:
        return

    # do the trace
    trace = KSegment(net=net, layer="F.Cu", point1=do.absolute_at, point2=di.absolute_at, width=TRACE_WIDTH_MM)
    board.add(trace)

"""
Connects the line to the MC
"""
def connect_line(board, mc, line, station_fp):
    if connect_line.gpio_ptr > len(MC_PADS["gpios"]):
        logging.critical(
            f"line {line.name} not connected! cannot have more than {
                len(MC_PADS['gpios'])} lines."
        )
        return

    gpio = MC_PADS["gpios"][connect_line.gpio_ptr]
    logging.info(f"connecting line {line.name} to gpio {gpio}")

    net = KNet(1000 + gpio, f"{line}")
    board.add(net)
    mc.pads[gpio].net = net
    station_fp.pads[LED_PADS["DI"]].net = net

    connect_line.gpio_ptr += 1
connect_line.gpio_ptr = 0


"""
Adds the main microcontroller
"""
def add_mc(board):
    mc = KFootprint(FOOTPRINT_MC, at=MC_POS)

    if not NO_POWER_NETS:
        mc.pads[MC_PADS["GND"]].net = NET_GND
        mc.pads[MC_PADS["VCC"]].net = NET_VCC

    board.add(mc)
    return mc


"""
adds labels for the lines to the next position in the TEXT_POS_MAPPING list
from the position of the text label of the station (just a best estimate)
"""
def add_line_label(board, line, station):
    text_pos_x = station.pos.x + TEXT_SPACING_MM * math.sin(-0.25*station.angle)
    text_pos_y = station.pos.y + TEXT_SPACING_MM * math.cos(-0.25*station.angle)

    text_pos = Vector(text_pos_x, text_pos_y)
    text = KGrText(
        layer="F.SilkS", text=f"{line.name}", at=text_pos
    )
    #text.font = LINE_FONT

    board.add(text)


"""
adds a board title top middle
"""
def add_title_label(board, map):
    if NO_TEXT:
       return

    logging.debug(f"add map title: {map.title}")
    text = KGrText(
        layer="F.SilkS", text=f"{map.title}", at=Point(map.width / 2, MARGIN["top"])
    )
    text.font = TITLE_FONT

    board.add(text)


"""
main pcb generation loop
"""
def generate_pcb(map):
    board = KProject(PCB_FILE)
    board.add(NET_GND)
    board.add(NET_VCC)

    add_title_label(board, map)

    mc = add_mc(board)

    for line in map.lines:
        prev_station = None
        prev_led = None
        if len(line.stations) == 0:
            continue
        for station in line.stations:
            if prev_station is not None:
                station.angle = calculate_angle(prev_station, station)
            else:
                station.angle = calculate_angle(station, line.stations[1])
            logging.debug(f"angle between stations: {station.angle}")
            led = create_station_footprint(station)
            board.add(led)
            cap = create_capacity(station)
            board.add(cap)

            if prev_station is not None:
                connect_stations(board, prev_station, station, prev_led, led)
            else:
                logging.debug(f"add line label {line.name} to start of line")
                add_line_label(board, line, station)
                connect_line(board, mc, line, led)

            prev_station = station
            prev_led = led

        logging.debug(f"add line label {line.name} to end of line")
        add_line_label(board, line, prev_station)

    board.save()
    return board


def main():
    parse_arguments()

    logging.basicConfig(
        level=LOG_LEVEL,
        format="%(asctime)s - %(levelname)s - %(message)s",
        stream=sys.stdout,
    )
    logging.info(f"creating {PCB_FILE} from {INPUT_XML}")

    start_time = time.perf_counter()
    map = parse_map(INPUT_XML)

    pcb = generate_pcb(map)

    end_time = time.perf_counter()
    elapsed = end_time - start_time
    logging.info(f"finished in {elapsed:.4f} seconds")


def parse_arguments():
    global SCALE, TEXT_SPACING_MM, NO_TEXT, NO_POWER_NETS, NO_TRACES, PCB_FILE, LOG_LEVEL, INPUT_XML

    parser = argparse.ArgumentParser(
        description="train map pcb generator",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Positional Argument
    parser.add_argument("input", help="input xml file")

    # Optional Arguments
    parser.add_argument("--scale", type=float, default=1.0,
                        help="scaling of the pcb")
    parser.add_argument(
        "--text-space",
        type=int,
        default=5,
        help="spacing of the text from the center point of led (mm)",
    )
    parser.add_argument(
        "--no-text",
        action="store_true",
        help="do not produce any text",
    )
    parser.add_argument(
        "--no-power-nets",
        action="store_true",
        help="do not connect GND and VCC (easier visual debugging)",
    )
    parser.add_argument(
        "--no-traces",
        action="store_true",
        help="do not implement traces between leds",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=str,
        default="trainmap.kicad_pcb",
        help="output file. should end in .kicad_pcb",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="count",
        default=0,
        help="verbose output (use -v, -vv, etc.)",
    )

    args = parser.parse_args()

    INPUT_XML = args.input
    SCALE = args.scale
    TEXT_SPACING_MM = args.text_space
    NO_TEXT = args.no_text
    NO_POWER_NETS = args.no_power_nets
    NO_TRACES = args.no_traces
    PCB_FILE = args.output

    if args.verbose > 0:
        LOG_LEVEL = logging.DEBUG

    # Basic Validation
    if not PCB_FILE.endswith(".kicad_pcb"):
        print(f"Warning: Output file '{
              PCB_FILE}' does not have .kicad_pcb extension.")

    if NO_TEXT:
        print("Mode: No text")

    if args.no_power_nets:
        print("Mode: Power nets disabled.")

    if NO_TRACES:
        print("Mode: No traces")


if __name__ == "__main__":
    main()