import os
from typing import List, Union
import logging
import re
from math import isclose

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s - %(levelname)s: %(message)s')

FLOAT_OR_INT = r'-?\d+(\.\d*)?'

RE_ENGINE_VERSION = re.compile(r'SteamEngine (\d+\.\d+\.\d+)')
RE_ENGINE_V5 = re.compile(r'\n(.*SteamEngine main.*)\n')
RE_LAYER_COUNT = re.compile(r';LAYER_COUNT:(\d+)', re.MULTILINE)
RE_LAYER_NUMBER = re.compile(r';LAYER:(\d+)')
RE_FOOTER_START = re.compile(r'(;TIME_ELAPSED:\d+\.\d+)\n[^;]', re.MULTILINE)
RE_FAN_SPEED = re.compile(r'(M106 .+)\n', re.MULTILINE)
RE_LAYER_ZHEIGHT = re.compile(r'G0 .*Z({})'.format(FLOAT_OR_INT))
RE_EXTRUSION = re.compile(r'G1 .*E({})'.format(FLOAT_OR_INT))
RE_MILLIMETERS = re.compile(r'({})mm'.format(FLOAT_OR_INT))

RE_LAYER_V3 = re.compile(r'(;LAYER:\d+)', re.MULTILINE)
RE_LAYER_V4 = re.compile(r'^G0.* Z\d+\.?\d*$|^;LAYER:0$', re.MULTILINE)

class GcodeFileAnalyzer:

    @classmethod
    def LoadFile(cls, path):
        logging.info('Loading %s', path)
        with open(path, 'r') as f:
            return cls([l.strip('\n') for l in f])

    @classmethod
    def SplitFileAtLayers(cls, path, layerNumbers: List[Union[int, str]]):
        gcode = cls.LoadFile(path)
        for i, lines in enumerate(gcode.ExportSplitted(layerNumbers)):
            outPath = os.path.join(os.path.dirname(path),
                                   '{:02d}_{}'.format(i, os.path.basename(path)))
            logging.info('Writing %s', outPath)
            with open(outPath, 'wt') as f:
                f.write('\n'.join(lines))

    def __init__(self, lines):
        self.curaVersion: tuple[int] = None
        self.lines: list[str] = lines
        self.header: list[str] = []
        self.footer: list[str] = []
        self.layers: list['Layer'] = []
        self.layerCount: int = 0
        self.fanSpeedCode: str = ''
        self.__AnalyzeGcode()

    def __AnalyzeGcode(self) -> None:
        logging.info('Analyzing GCode')
        fullCode = '\n'.join(self.lines)

        versionMatch = RE_ENGINE_VERSION.search(fullCode)
        if versionMatch:
            self.curaVersion = tuple(int(v) for v in versionMatch.group(1).split('.'))
            logging.info(f'Cura version {self.curaVersion}')
        else:
            v5engineMatch = RE_ENGINE_V5.search(fullCode)
            if v5engineMatch:
                logging.info('Cura version not specified. Assuming 5.x')
                logging.info(f' -> {v5engineMatch.group(1)}')
                self.curaVersion = (5, 0, 0)
            else:
                raise ValueError('Could not find ";Generated with Cura_SteamEngine" '
                                 'in GCODE. Can\'t proceed')

        # Find footer
        footerStartMatch = RE_FOOTER_START.search(fullCode)
        startOfFooterIndex = self.lines.index(footerStartMatch.group(1)) + 1
        self.footer = self.lines[startOfFooterIndex:]

        layersIndices = self.__FindLayers(fullCode, startOfFooterIndex)

        # Find header
        self.header = self.lines[:layersIndices[0]]

        # Find fan speed code
        fanSpeedMatch = RE_FAN_SPEED.search(fullCode)
        if fanSpeedMatch:
            self.fanSpeedCode = fanSpeedMatch.group(1)

    def __FindLayers(self, fullCode, startOfFooterIndex):
        # Find layers
        layerCountMatch = RE_LAYER_COUNT.search(fullCode)
        self.layerCount = int(layerCountMatch.group(1))

        if self.curaVersion < (4, 0, 0):
            # Layers start at line ;LAYER:<layer_no>
            layersList = RE_LAYER_V3.findall(fullCode)

        else:  # Version > 4.0.0
            # Layers start already before line ;LAYER:<layer_no>
            # (except for ;LAYER:0)
            # Look for actual Z-position changes
            layersList = RE_LAYER_V4.findall(fullCode)
            layersList.remove(layersList[1])  # drop Z-change for layer 0, because it happens after the ;LAYER:0 line

        if len(layersList) != self.layerCount:
            logging.warning('Gcode states wrong layer count {}. But found {} layers.'
                        .format(layerCountMatch.group(0), len(layersList)))
            self.layerCount = len(layersList)

        layersIndices = [self.lines.index(layer) for layer in layersList] + [startOfFooterIndex]
        self.layers = [Layer(self.lines[i:j]) for i, j in zip(layersIndices[:-1], layersIndices[1:])]
        return layersIndices

    def ExportSplitted(self, splitBeforeLayers: list) -> list:
        splitBeforeLayers = self._FindLayerIndices(splitBeforeLayers)
        assert all(l in range(1, self.layerCount) for l in splitBeforeLayers)
        outGcodes = []
        for startLayerIndex, endLayerIndex in zip([0] + splitBeforeLayers,
                                                  splitBeforeLayers + [self.layerCount]):
            startLayer = self.layers[startLayerIndex]
            endLayer = self.layers[endLayerIndex - 1]
            outGcodes.append(self._ExtractLayerSubset(startLayer, endLayer))
        return outGcodes

    def _FindLayerIndices(self, layerIndicesOrHeights: List[Union[int, str]]) -> List[int]:
        """ Translate heights to the corresponding layer index
        """
        layerIndices = []
        for layerIndexOrHeight in layerIndicesOrHeights:
            if isinstance(layerIndexOrHeight, int):
                # already a layer index
                # minus one because first layer input starts at one (similar to Cura)
                layerIndices.append(layerIndexOrHeight - 1)
            elif isinstance(layerIndexOrHeight, str):
                # extract millimeters
                millisMatch = RE_MILLIMETERS.search(layerIndexOrHeight)
                if not millisMatch:
                    raise ValueError('{} is not a valid millimeters measure'.format(layerIndexOrHeight))
                # find layer at index
                height: float = float(millisMatch.group(1))
                for layer in self.layers:
                    if isclose(height, layer.zHeight, abs_tol=0.001):
                        layerIndices.append(layer.number)
                        break
                else:
                    raise RuntimeError('Layer at Z position {} could not be found'.format(layerIndexOrHeight))
        return layerIndices

    def _ExtractLayerSubset(self, startLayer: 'Layer', endLayer: 'Layer') -> List[str]:
        # Add default header
        lines = list(self.header)
        logging.info('Exporting Layers {} ({} mm) ... {} ({} mm)'
                     .format(startLayer.number + 1, startLayer.zHeight,
                             endLayer.number + 1, endLayer.zHeight))

        if startLayer.number != 0:
            # Add extra initalization
            lines += self._GetAdditionalHeader(startLayer)
        # Add layers
        for layer in self.layers[startLayer.number:endLayer.number + 1]:
            lines += layer.lines
        # Add footer
        if endLayer.number + 1 == self.layerCount:
            lines += self.footer  # default footer
        else:
            # End extrusion must be patched
            lines += self._GetPatchedFooter(endLayer)
        return lines

    def _GetAdditionalHeader(self, fromLayer: 'Layer') -> List[str]:
        return [
            '', ';GENERATED_EXTRA_INIT',
            self.fanSpeedCode,
            'G0 Z{:.3f}'.format(fromLayer.zHeight + 1),  # Approach from above
            'G92 E{:.3f}'.format(fromLayer.startExtrusion - 10),  # Extrude 10 more mm
            '',
        ]

    def _GetPatchedFooter(self, endLayer: 'Layer') -> List[str]:
        footerExtrusionIndex = -1
        for i, line in enumerate(self.footer):
            if RE_EXTRUSION.search(line):
                footerExtrusionIndex = i
                break
        else:
            raise RuntimeError('Extrusion not found in footer')

        patchedFooter = list(self.footer)
        patchedFooter[footerExtrusionIndex] = '; {}'.format(patchedFooter[footerExtrusionIndex])
        return [
            '', ';PATCHED_FOOTER',
            'G1 F1500 E{:.3f}'.format(endLayer.endExtrusion - 6.5),
            '',
        ] + patchedFooter


class Layer:
    def __init__(self, lines):
        self.lines = lines
        self.zHeight: float = 0
        self.startExtrusion: float = 0
        self.endExtrusion: float = 0
        self.number: int = 0
        self.__AnalyzeLayerGcode()

    def __AnalyzeLayerGcode(self):
        """ Analyze layer GCode to find its Z height and extrusion position
        """
        # Find Layer number
        for line in self.lines:
            match = RE_LAYER_NUMBER.search(line)
            if match:
                self.number = int(match.group(1))
                break
        else:
            raise RuntimeError('No layer number found for {}'.format(self.lines[0]))

        # Find Z height
        for line in self.lines:
            match = RE_LAYER_ZHEIGHT.search(line)
            if match:
                self.zHeight = float(match.group(1))
                break
        else:
            raise RuntimeError('No Z height found for {}'.format(self.lines[0]))

        # Find Start Extrusion
        for line in self.lines:
            match = RE_EXTRUSION.search(line)
            if match:
                self.startExtrusion = float(match.group(1))
                break
        else:
            raise RuntimeError('No startExtrusion found for {}'.format(self.lines[0]))

        # Find End Extrusion
        for line in reversed(self.lines):
            match = RE_EXTRUSION.search(line)
            if match:
                self.endExtrusion = float(match.group(1))
                break
        else:
            raise RuntimeError('No endExtrusion found for {}'.format(self.lines[0]))


if __name__ == '__main__':
    GcodeFileAnalyzer.SplitFileAtLayers('test/BigL.gcode', ['2.8mm', 300])