import re
import hashlib
import zipfile
from pathlib import Path
from functools import partial, cache
from dataclasses import dataclass
from pymoku.parts import (
    log, link, copy, BuildStep, ProcessStep, SourceFile, ParseNode,
    global_alias, arg)


show_stats = arg('--stats', action='store_true')


_with_rich = True
try:
    from rich.table import Table
    from rich.bar import Bar
    from rich.text import Text
    from rich.columns import Columns
except ImportError:
    _with_rich = False


class VivadoStep(ProcessStep):
    def __init__(self, script, *args, **kwargs):
        self.dcp = kwargs.pop('dcp', None)
        self.tcl_args = kwargs.pop('tcl_args', ())
        super().__init__(**kwargs)
        # resolve script to an actual path
        self.script = list(self.depends_on(link(script, self.workdir)))[0].path

    def parse_tree(self):
        fn = self.update_progress
        start_node = ParseNode()
        node = start_node
        node = node.add(r'source.*\.tcl', partial(fn, completed=2, phase='running'))
        node = node.add(r'INFO:.*Exiting Vivado', partial(fn, completed=99))
        return start_node

    def process_log(self, logline):
        line = logline.decode('utf8').strip()
        if re.match(r'WARNING', line):
            self.log.warning(line)
        elif re.match(r'CRITICAL WARNING', line):
            self.log.error(line)
        elif re.match(r'ERROR', line):
            self.log.error(line)
        return super().process_log(logline)

    def _arg_strings(self):
        """ Convert args to string to pass to subprocess
        """
        for a in self.tcl_args:
            if isinstance(a, SourceFile):
                a = self.relative_to(self.config_path / a.path)  # TODO figure out this config_path
            yield str(a)

    def command(self):
        cmd = f'vivado -mode batch -source {self.relative_to(self.script)}'
        cmd += f' {self.dcp}' if self.dcp else ''
        if self.tcl_args:
            cmd += f' -tclargs {" ".join(self._arg_strings())}'
        return cmd


class ProcessorStep(VivadoStep):
    def parse_tree(self):
        example_log = [
            (r'source.*\.tcl', dict(completed=2)),
            (r'.*create_root_design', dict(completed=12, phase='running')),
            (r'VHDL Output written to', dict(completed=13)),
            (r'Starting synth_design', dict(completed=21)),
            (r'Starting RTL Elaboration', dict(completed=25)),
            (r'Start Handling Custom Attributes', dict(completed=32)),
            (r'Processing XDC Constraints', dict(completed=36)),
            (r'Completed Processing XDC Constraints', dict(completed=39)),
            (r'Finished Constraint Validation', dict(completed=41)),
            (r'Finished RTL Optimization Phase 2', dict(completed=50)),
            (r'Start RTL Component Statistics', dict(completed=52)),
            (r'Start Applying XDC Timing Constraints', dict(completed=65)),
            (r'Start Timing Optimization', dict(completed=66)),
            (r'Start Technology Mapping', dict(completed=73)),
            (r'Start IO Insertion', dict(completed=81)),
            (r'Start Final Netlist Cleanup', dict(completed=82)),
            (r'Finished Final Netlist Cleanup', dict(completed=84)),
            (r'Start Renaming Generated Instances', dict(completed=85)),
            (r'Start Rebuilding User Hierarchy', dict(completed=86)),
            (r'Start Renaming Generated Ports', dict(completed=87)),
            (r'synth_design completed successfully', dict(completed=95)),
            (r'INFO:.*Exiting Vivado', dict(completed=99)),
        ]

        start_node = ParseNode()
        fn = self.update_progress
        node = start_node
        for msg, kwargs in example_log:
            node = node.add(msg, partial(fn, **kwargs))

        return start_node


class SynthStep(VivadoStep):
    def parse_tree(self):
        example_log = [
            (r'source.*\.tcl', dict(completed=1)),
            (r'Starting synth_design', dict(completed=3, phase='running')),
            (r'Starting RTL Elaboration', dict(completed=4)),
            (r'Start Handling Custom Attributes', dict(completed=6)),
            (r'Netlist sorting complete', dict(completed=7)),
            (r'Processing XDC Constraints', dict(completed=10)),
            (r'Completed Processing XDC Constraints', dict(completed=12)),
            (r'Finished Constraint Validation', dict(completed=14)),
            (r'Finished RTL Optimization Phase 2', dict(completed=30)),
            (r'Start RTL Component Statistics', dict(completed=33)),
            (r'Finished RTL Component Statistics', dict(completed=33)),
            (r'Start Part Resource Summary', dict(completed=33)),
            (r'Finished Part Resource Summary', dict(completed=33)),
            (r'Start Cross Boundary and Area Optimization', dict(completed=33)),
            (r'Finished Cross Boundary and Area Optimization', dict(completed=47)),
            (r'Start Applying XDC Timing Constraints', dict(completed=48)),
            (r'Start Timing Optimization', dict(completed=50)),
            (r'Start Technology Mapping', dict(completed=58)),
            (r'Start IO Insertion', dict(completed=70)),
            (r'Start Final Netlist Cleanup', dict(completed=70)),
            (r'Finished Final Netlist Cleanup', dict(completed=73)),
            (r'Finished IO Insertion', dict(completed=74)),
            (r'Start Renaming Generated Instances', dict(completed=74)),
            (r'Start Handling Custom Attributes', dict(completed=78)),
            (r'Start Writing Synthesis Report', dict(completed=78)),
            (r'Synthesis Optimization Runtime', dict(completed=79)),
            (r'Netlist sorting complete', dict(completed=83)),
            (r'Reading XDEF placement', dict(completed=92)),
            (r'synth_design completed successfully', dict(completed=93)),
            (r'Writing XDEF routing', dict(completed=95)),
            (r'INFO:.*Exiting Vivado', dict(completed=99)),
        ]

        start_node = ParseNode()
        fn = self.update_progress
        node = start_node
        for msg, kwargs in example_log:
            node = node.add(msg, partial(fn, **kwargs))

        return start_node


class RouteStep(VivadoStep):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.wns = 0.0
        self.tns = 0.0
        self.phase = ''
        self.overlaps = ''

    def stats(self, match):
        self.wns = float(match.group('wns'))
        self.tns = float(match.group('tns'))
        log.debug(f'{self} WNS={self.wns:6.3f} TNS={self.tns:6.3f}')
        col = '[red]' if self.tns < 0.0 else ''
        self.update_progress(comment=f'{col}WNS={self.wns:6.3f} TNS={self.tns:6.3f}')

    def match_phase(self, match):
        self.phase = match.group('phase')
        log.debug(f'{self} Phase: {self.phase}')

    def match_overlaps(self, match):
        self.overlaps = match.group('overlaps')
        log.debug(f'{self} Overlaps: {self.overlaps}')
        self.update_progress(comment=f'WNS={self.wns:6.3f} TNS={self.tns:6.3f} {self.overlaps}')

    def process_log(self, logline):
        line = logline.decode('utf8')

        if (match := re.match(r'.*WNS=(?P<wns>[-\d]+\.[\d]+)[\s\|]*TNS=(?P<tns>[-\d]+\.[\d]+)', line)) is not None:
            self.stats(match)

        if (match := re.match(r'Phase (?P<phase>[\d]+[\.\d]*)', line)) is not None:
            self.match_phase(match)

        if (match := re.match(r'\sNumber of Nodes with overlaps = (?P<overlaps>[\d]+)', line)) is not None:
            self.match_overlaps(match)

        return super().process_log(logline)

    def parse_tree(self):
        example_log = [
            (r'Command: open_checkpoint', dict(completed=0, phase='open_checkpoint')),
            (r'Command: opt_design', dict(completed=3, phase='opt_design')),
            (r'Phase 1 Initialization', dict(completed=4)),
            (r'Phase 1.2 Setup Constraints And Sort Netlist', dict(completed=9)),
            (r'Phase 6 BUFG optimization', dict(completed=10)),
            (r'Command: place_design', dict(completed=10, phase='place_design')),
            (r'Command: phys_opt_design', dict(completed=30, phase='phys_opt_design')),
            (r'Command: route_design', dict(completed=50, phase='route_design')),
            (r'Command: phys_opt_design', dict(completed=90, phase='phys_opt_design')),
        ]

        start_node = ParseNode(strict=True)
        fn = self.update_progress
        node = start_node
        for msg, kwargs in example_log:
            node = node.add(msg, partial(fn, **kwargs))

        return start_node


# class GenIPStep(VivadoStep):
#     def __init__(self, xci):
#         self.xci = xci
#         name = f'{xci.path.stem}'
#         super().__init__('scripts/ip.tcl', xci.path.name, name=name)
#         self.depends_on(Copy(xci, XCIFile(self.workdir / xci.path.name)))

#     # def workdir(self):
#     #     return self.xci.outputs()[0].path.parent

#     def outputs(self):
#         return [DCPFile(self.workdir / self.xci.path.with_suffix('.dcp'))]

#     def parse_tree(self):
#         fn = self.update_progress
#         start_node = ParseNode()
#         node = start_node.add(r'source.*ip\.tcl', partial(fn, completed=5, phase='starting'))
#         node = node.add(r'Starting synth_design', partial(fn, completed=10, phase='synth'))
#         node = node.add(r'Starting RTL Elaboration', partial(fn, completed=30, phase='elaborate'))
#         node = node.add(r'Finished RTL Elaboration', partial(fn, completed=60, phase='finishing'))
#         return start_node


# @expects('support-processor-tcl')
def processor(env):
    proc = ProcessorStep('support-processor-tcl', name='processor', extra_env=env)
    proc.depends_on(link('tcl/processor_bd.tcl', proc.workdir))
    proc.creates('processor.dcp')
    proc.creates('.gen/sources_1/bd/processor/hdl/processor_wrapper.vhd')
    return proc


class BitbinStep(ProcessStep):
    def __init__(self, source, platform, **kwargs):
        self.arch = platform.arch
        self._workdir = source.path.parent
        super().__init__()
        self.depends_on(source)

        self.bitfile = source.path
        binfile = self.bitfile.with_suffix('.bit.bin')
        self.creates(Bitstream(binfile, platform=platform, **kwargs))

    @property
    def workdir(self):
        return self._workdir

    def command(self):
        bs = self.bitfile.name
        return f"${{BOOTGEN:-`which bootgen`}} -w on -image {bs}.bif -arch {self.arch} -process_bitstream bin"

    async def do_step(self):
        biffile = Path(self.bitfile).with_suffix('.bit.bif')
        biffile.write_text(f"all: {{ {self.bitfile.name} }}")
        await super().do_step()


class AbstractShell(VivadoStep):
    def __init__(self, platform, cell_name, pblock_name):
        self.platform = platform
        self.cell_name = cell_name
        self.pblock_name = pblock_name
        name = f'{platform.name}-{self.safe_cell_name}'
        shell_dcp = f'{name}.dcp'

        super().__init__(f'support-shell-{platform.arch}-tcl',
                         name=name,
                         dcp='platform_route.dcp',
                         tcl_args=(cell_name, shell_dcp))

        self.depends_on(link(platform.route(), self.workdir))
        self.creates(shell_dcp)

    @property
    def plat_name(self):
        return self.platform.name

    @property
    def path(self):
        return next(iter(self.outputs())).path

    @property
    def safe_cell_name(self):
        name = self.cell_name.translate(
            str.maketrans({"[": "_", "]": None, "*": "_"}))
        name = name.replace('_INSTR', '')
        return name

    @property
    def slot(self):
        """ This makes some big assumptions about cell naming conventions"""
        if (m := re.search(r'INSTRS\[(\d+)\]', self.cell_name)) is None:
            return None
        return int(m.group(1))

    @cache
    def sha(self):
        sha = hashlib.sha256(self.path.read_bytes())
        return sha.hexdigest()


class Platform(BuildStep):
    def __init__(self, env, ip, vhdl, slots, slot_ains, slot_aouts, slot_cbufs):
        self.slots = slots
        self.slot_ains = slot_ains
        self.slot_aouts = slot_aouts
        self.slot_cbufs = slot_cbufs
        self.env = env
        self.ip = ip
        self.vhdl = vhdl
        super().__init__(name=f'{self.hgboard}-{slots:d}')
        self._shells = {}

        self.depends_on(self.bitstreams())
        self.depends_on(Timing(self.reports(), self.name))

    @property
    def hgboard(self):
        return self.env['HG_BOARD']

    @property
    def part(self):
        # TODO remove this env stuff
        return self.env['HG_PART']

    @property
    def hwver(self):
        return self.env['HG_HWVER']

    @property
    def arch(self):
        return self.env['HG_ARCH']

    @property
    def top(self):
        return self.env['HG_TOP']

    @property
    def sha(self):
        return self._platform_bit.sha()

    def outputs(self):
        return self.inputs()

    def slot_shells(self):
        return sorted([v for v in self._shells.values() if v.slot is not None], key=lambda s: s.slot)

    @cache
    def create_shell(self, cell_name, pblock_name):
        if cell_name in self._shells:
            return self._shells[cell_name]

        shell = AbstractShell(self, cell_name, pblock_name)
        self._shells[cell_name] = shell
        return shell

    def shell(self, cell_name):
        return self._shells[cell_name]

    @cache  # multiple calls should not create new instances
    def synth(self):
        synth = SynthStep('support-synth-tcl',
                          name=f'synth-{self.slots:d}',
                          tcl_args=(self.slots, self.slot_ains, self.slot_aouts, self.slot_cbufs),
                          extra_env=self.env)
        synth.depends_on(link('processor', synth.workdir),
                         link('constraints/constraints.xdc', synth.workdir / 'constraints'),
                         link(self.vhdl, synth.workdir / 'vhdl'),
                         link('support-vhdl', synth.workdir / 'vhdl/support'),
                         link(self.ip, synth.workdir / 'ip'))
        synth.creates('platform_synth.dcp')
        global_alias('all-synth', synth)
        return synth

    @cache
    def route(self):
        synth = self.synth()
        route = RouteStep('support-route-tcl',
                          name=f'route-{self.slots:d}',
                          dcp='platform_synth.dcp')
        route.depends_on(link(synth, route.workdir),
                         link('constraints/constraints.xdc', route.workdir / 'constraints'),
                         link(f'constraints/floorplan_{self.slots:d}.xdc', route.workdir / 'constraints'),
                         link('support-slot-constraints', route.workdir))
        return route.creates('platform_route.dcp')

    @cache
    def bitstreams(self):
        route = self.route()
        bitstreams = VivadoStep('support-bitstreams-tcl',
                                name=f'bitstreams-{self.slots:d}',
                                dcp='platform_route.dcp')
        bitstreams.depends_on(link(route, bitstreams.workdir))
        self._platform_xsa = bitstreams.creates('platform.xsa')
        bitbin = BitbinStep(bitstreams.creates('platform.bit').pop(), platform=self)
        self._platform_bit = list(bitbin.outputs())[0]
        self._empty_bits = [
            BitbinStep(bitstreams.creates(
                f'platform_pblock_INSTR{i:d}_partial.bit').pop(),
                platform=self, slot=i)
            for i in range(self.slots)]
        return (self._platform_bit, self._empty_bits)

    @cache
    def reports(self):
        route = self.route()
        name = f'{self.hgboard}-{self.slots:d}'

        s = VivadoStep('support-platform-reports-tcl',
                       name=f'reports-{name}',
                       dcp='platform_route.dcp',
                       tcl_args=(name,))
        s.depends_on(link(route, s.workdir))
        s.creates(f'{name}_util.rpt',
                  f'{name}_timing.rpt',
                  f'{name}_util_hier.rpt')
        return s

    def platform_bit(self):
        return self._platform_bit

    def empty_bits(self):
        return self._empty_bits


class Bitstream(SourceFile):
    def __init__(self, path, *args, **kwargs):
        self.slot = kwargs.pop('slot', None)
        self.platform = kwargs.pop('platform')
        self.manifest_extra = kwargs.pop('manifest_extra', {})
        super().__init__(path, *args, **kwargs)

    @cache
    def sha(self):
        sha = hashlib.sha256(self.path.read_bytes())
        return sha.hexdigest()


class PythonZip(BuildStep):
    def __init__(self, outfile, **kwargs):
        super().__init__(**kwargs)
        self.path = next(iter(self.creates(outfile))).path
        self.platform = set()  # HACK
        self.depends_on('support-slot-python')

    async def do_step(self):
        with zipfile.ZipFile(self.path, mode='w') as f:
            # using PyZipFile causes issues across different python versions
            for inp in self.inputs():
                if inp.path.name == 'api_main.py':
                    f.write(inp.path, arcname='__main__.py')
                else:
                    f.write(inp.path, arcname='__pymoku_plugin/' + inp.path.name)  # TODO don't flatten everything

    @cache
    def sha(self):
        sha = hashlib.sha256(self.path.read_bytes())
        return sha.hexdigest()


class RMSynth(SynthStep):
    def __init__(self, top, vhdl, fpga_part, config_name=None, ip=[], generics={}, env={},
                 synth_tcl='support-slot-synth-tcl',
                 constraints='support-slot-constraints'):
        self.config_name = config_name
        self.generics = generics

        generics_str = ' '.join([f'{k}={str(v).lower()}' for k, v in generics.items()])
        generics_str = f'"{generics_str}"'
        super().__init__(synth_tcl,
                         name=f'{config_name}-synth',
                         extra_env=env,
                         tcl_args=(top, fpga_part, generics_str))

        for core in ip:
            self.depends_on(copy(core, self.workdir / 'ip' / core.path.stem))
        self.depends_on(link('support-slot-vhdl', self.workdir / 'vhdl/support')),
        self.depends_on(link(vhdl, self.workdir / 'vhdl'))
        self.depends_on(link(constraints, self.workdir / 'constraints'))
        self.creates('synth.dcp')
        global_alias('all-synth', self)


class RMRoute(RouteStep):
    def __init__(self, synth, shell,
                 route_tcl='support-slot-route-tcl',
                 constraints='support-slot-constraints'):
        self.config_name = f'{synth.config_name}-{shell.safe_cell_name}'
        self.shell = shell
        self.synth = synth
        super().__init__(route_tcl,
                         name=f'{self.config_name}-route',
                         dcp=shell.path.name,
                         tcl_args=(synth.workdir.absolute() / 'synth.dcp',
                                   shell.cell_name,
                                   'route.dcp'))
        self.depends_on(synth)
        self.depends_on(link(shell, self.workdir))
        self.depends_on(link(constraints, self.workdir / 'constraints'))
        self.creates('route.dcp')


class RMInstance(BuildStep):
    def __init__(self, route,
                 bitstream_tcl='support-slot-bitstream-tcl',
                 reports_tcl='support-slot-reports-tcl'):
        self.bitstream_tcl = bitstream_tcl
        self.reports_tcl = reports_tcl
        super().__init__()
        self.depends_on(self.reports(route))
        self.creates(self.bitstream(route).outputs())
        self.creates(self.reports(route).outputs())

    @cache
    def bitstream(self, route):
        b = VivadoStep(self.bitstream_tcl,
                       name=f'{route.config_name}-bitstream',
                       dcp='route.dcp',
                       tcl_args=(route.shell.cell_name, f'out_{route.shell.safe_cell_name}'))
        b.depends_on(link(route, b.workdir))
        bs = BitbinStep(
            b.creates(f'out_{route.shell.safe_cell_name}.bit').pop(),
            platform=route.shell.platform,
            slot=route.shell.slot,
            manifest_extra={
                k.lower(): v for k, v in route.synth.generics.items()
            })
        return bs

    @cache
    def reports(self, route):
        s = VivadoStep(self.reports_tcl,
                       name=f'{route.config_name}-reports',
                       dcp='route.dcp',
                       tcl_args=(route.shell.cell_name, route.shell.pblock_name))
        s.depends_on(link(route, s.workdir))
        s.creates('pblock_util.rpt', 'timing.rpt')

        name = f'{route.config_name}'
        if self.prefix:
            name = f'{self.prefix}-{name}'

        return Timing(s, name=name)


class Timing(BuildStep):
    results_table = None

    @dataclass
    class Stats(object):
        clb: tuple[int, int, float] = (0, 0, 0)
        dsp: tuple[int, int, float] = (0, 0, 0)
        bram: tuple[int, int, float] = (0, 0, 0)
        wns: float = 0

    def __init__(self, reports, name, **kwargs):
        super().__init__(**kwargs)
        self.row_name = name
        self.depends_on(reports)

    def state(self):
        return b''

    def current(self, in_state):
        return False

    async def do_step(self):
        if Timing.results_table is None and _with_rich:
            table = Table(show_edge=False, pad_edge=False, padding=0)
            Timing.results_table = table
            self.output_table.add_row(table)

            table.add_column(" ", justify="left", style="cyan", no_wrap=True)
            table.add_column("CLB", justify="center", width=17)
            table.add_column("DSP", justify="center", width=27)
            table.add_column("BRAM18", justify="center", width=27)
            table.add_column("WNS", justify="center", width=6)

        data = {}
        for inp in self.inputs():
            if inp.path.name.endswith('util.rpt'):
                data[self.row_name] = data.get(self.row_name, self.Stats())
                self.parse_util(inp.path, data[self.row_name])
            elif inp.path.name.endswith('timing.rpt'):
                data[self.row_name] = data.get(self.row_name, self.Stats())
                self.parse_timing(inp.path, data[self.row_name])
            else:
                continue

        for k, v in data.items():
            self.print_row(k, v)

    def print_row(self, name, stats):
        if _with_rich and (show_stats or stats.wns < 0.0):
            text = [Text(name),
                    Columns([Text(f'{stats.clb[2]:5.1f}%'),
                             Bar(100.0, 0.0, stats.clb[2], width=10, bgcolor='grey11', color='sandy_brown' if stats.clb[2] > 80.0 else 'default')]),
                    Columns([Text(f'{stats.dsp[2]:5.1f}%'),
                             Bar(100.0, 0.0, stats.dsp[2], width=10, bgcolor='grey11'),
                             Text(f'{stats.dsp[0]:4d}/{stats.dsp[1]:4d}')]),
                    Columns([Text(f'{stats.bram[2]:5.1f}%'),
                             Bar(100.0, 0.0, stats.bram[2], width=10, bgcolor='grey11'),
                             Text(f'{stats.bram[0]:4d}/{stats.bram[1]:4d}')]),
                    Text(f'{stats.wns:6.3f}', style='red' if stats.wns < 0.0 else None)]
            self.results_table.add_row(*text)

        else:
            self.log.info(f'{name}:'
                          f' CLB {stats.clb[2]:5.1f}%'
                          f' DSP {stats.dsp[2]:5.1f}%'
                          f' BRAM {stats.bram[2]:5.1f}%'
                          f' WNS {stats.wns:6.3f}')
            return

    def parse_util(self, path, stats):
        rpt = path.read_text()

        m = re.search(r'^\|\s(?:CLB|Slice)\s+\|.*\|$', rpt, flags=re.MULTILINE)
        row = [x.strip() for x in m.group(0).split('|')]
        used, _, _, avail, util = row[-6:-1]
        stats.clb = int(used), int(avail), float(util)

        m = re.search(r'^\|\sDSPs\s+\|.*\|$', rpt, flags=re.MULTILINE)
        row = [x.strip() for x in m.group(0).split('|')]
        used, _, _, avail, util = row[-6:-1]
        stats.dsp = int(used), int(avail), float(util)

        m = re.search(r'^\|\sBlock RAM Tile\s+\|.*\|$', rpt, flags=re.MULTILINE)
        row = [x.strip() for x in m.group(0).split('|')]
        used, _, _, avail, util = row[-6:-1]
        stats.bram = int(2 * float(used)), int(avail), float(util)

    def parse_timing(self, path, stats):
        rpt = path.read_text()
        # Slack (MET) :             0.140ns  (arrival time - required time)
        m = re.search(r'\(\bVIOLATED\b\)\s\:\s+([\-\d\.]+)ns', rpt, flags=re.MULTILINE)

        if m is not None:
            stats.wns = float(m.group(1))


# TODO ip is currently synth'd in each individual step that uses it which is
# a bit redundant.  There are some issues around vivado outputting encrypted
# netlists (.edn), as well as our ability to reuse a single XCI for multiple
# platforms
def gen_ip(xcis, platform):
    steps = []
    for xci in xcis:
        assert isinstance(xci, SourceFile)
        s = VivadoStep('support-synthip-tcl',
                       name=f'ip-{platform.hgboard}-{xci.path.stem}',
                       tcl_args=(xci.path.name,),
                       extra_env=platform.env)
        s.depends_on(link(xci, s.workdir))
        if (coe := xci.path.with_suffix('.coe')).exists():
            s.depends_on(link(coe, s.workdir))

        s.creates(f'{xci.path.stem}.dcp')
        steps.append(s)
    return steps
