#!/bin/python
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,
)
import xml.etree.ElementTree as ET

# Constants / Global Variables set through argparse
SCALE = 1.0
TEXT_SPACING_MM = 12
NO_POWER_NETS = False
PCB_FILE = "trainmap.kicad_pcb"
LOG_LEVEL = logging.INFO
INPUT_XML = ""
MC_POS = Point(0, 0)
TITLE_FONTFACE = ""
FONTFACE = ""
MARGIN = {"left": 10, "top": 5}

# Fonts
LINE_FONT = Font(size=(16, 16))
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"

# Mappings
TEXT_POS_MAPPING = {
    "above": Vector(0, TEXT_SPACING_MM),
    "right": Vector(TEXT_SPACING_MM, 0),
    "below": Vector(0, -TEXT_SPACING_MM),
    "left": Vector(-TEXT_SPACING_MM, 0),
}
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 = root.find("title").text
    map = Map(title)
    logging.debug(f"map title: {title}")

    width = 0
    height = 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 = int(station.get("id"))
            name = station.text
            x = int(station.get("x"))
            # search width
            if x > width:
                width = x
            y = int(station.get("y"))
            # search height
            if y > height:
                height = y
            text_pos = station.get("textpos")
            text_angle = int(station.get("angle"))

            station_parsed = Station(id, name, x, y, text_angle, text_pos)
            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


def create_capacity(station):
    cap_pos_index = ([*TEXT_POS_MAPPING].index(station.text_pos) + 2) % len(
        TEXT_POS_MAPPING
    )
    cap_pos = [*TEXT_POS_MAPPING.values()][cap_pos_index]
    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


"""
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)

    # text
    ref = fp.properties[KFootprintProperty.Type.Reference]
    ref.value = station.name
    ref.angle = station.angle
    ref.at = TEXT_POS_MAPPING[station.text_pos]

    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)
    fp1.pads[LED_PADS["DO"]].net = net
    fp2.pads[LED_PADS["DI"]].net = net


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):

    # get next position
    text_pos_index = ([*TEXT_POS_MAPPING].index(station.text_pos) + 1) % len(
        TEXT_POS_MAPPING
    )
    text_pos = [*TEXT_POS_MAPPING.values()][text_pos_index]
    line_pos = Point(station.pos.x + 2 * text_pos.x,
                     station.pos.y + 2 * text_pos.y)

    text = KGrText(
        layer="F.SilkS", text=f"{line.name}", at=line_pos, angle=station.angle
    )
    text.font = LINE_FONT

    board.add(text)


"""
adds a board title top middle
"""


def add_title_label(board, map):
    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
        for station in line.stations:
            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_POWER_NETS, 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-power-nets",
        action="store_true",
        help="do not connect  (easier visual debugging)",
    )
    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_POWER_NETS = args.no_power_nets
    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 args.no_power_nets:
        print("Mode: Power nets disabled.")


if __name__ == "__main__":
    main()