#!/bin/env python3
import os
import ast
import argparse
import io
import sys
import re
from pathlib import Path
import yaml
import graphviz


def log(msg):
    """Log message if in debug mode"""
    if args.debug:
        print(msg)


class TestStepVisitor(ast.NodeVisitor):
    """
    A custom test step visitor to grab the test description (docstring)
    and the test steps test.step(.....) for the test case.

    """
    def __init__(self):
        self.name = ""
        self.description = ""
        self.test_steps = []

    def visit_Module(self, node):
        """Extract docstring from a test."""
        docstring = ast.get_docstring(node)
        if docstring:
            lines = docstring.splitlines()
            newline_parsed = False

            for line in lines:
                if self.name == "":
                    self.name = line.strip()
                elif newline_parsed is False and line == "":
                    newline_parsed = True
                    continue  # Skip mandatory newline
                else:
                    self.description = f"{self.description}\n{line}"
            self.description = self.description.strip()
        self.generic_visit(node)

    def _rebuild_fstring(self, node: ast.JoinedStr) -> str:
        """
        Rebuild an f-string (JoinedStr) by concatenating literal pieces and placeholders.
        """
        parts = []
        for value in node.values:
            if isinstance(value, ast.Constant):
                parts.append(value.value)
            elif isinstance(value, ast.FormattedValue):
                expression = ast.unparse(value.value) if hasattr(ast, "unparse") else ast.dump(value.value)
                parts.append(f"{{{expression}}}")
        return "".join(parts)

    # Check for test.step() for the actual test steps
    def visit_Call(self, node):
        """Extract test.step from test"""
        if isinstance(node.func, ast.Attribute) and node.func.attr == 'step':
            arg = node.args[0]
            if isinstance(node.func.value, ast.Name) and node.func.value.id == 'test':
                if node.args:
                    if isinstance(arg, ast.Constant):
                        self.test_steps.append(node.args[0].value)
                    elif isinstance(arg, ast.JoinedStr):
                        fstring_value = self._rebuild_fstring(arg)
                        self.test_steps.append(fstring_value)
        self.generic_visit(node)  # Continue visiting other nodes


class TestCase:
    """Test specifcation resources for a test case"""
    def __init__(self, directory, name="test.py", rootdir=None):
        self.test_dir = Path(directory)
        if rootdir:
            rootdir = Path(f"{rootdir}")
            self.test_dir = self.test_dir.relative_to(rootdir)
        self.topo_dot = f"{directory}/topology.dot"
        self.topo_img = f"{directory}/topology"
        self.test_path = f"{directory}/{name}"
        self.spec_path = f"{directory}/{os.path.splitext(name)[0]}.adoc"
        log(f"TestCase({self.test_path})")

    def gen_topology(self):
        """Generate SVG file from the topology.dot file"""
        with open(self.topo_dot, 'r', encoding='utf-8') as dot_file:
            svg_file = f"{self.topo_img}.svg"

            try:
                dot_graph = dot_file.read()
                graph = graphviz.Source(dot_graph)
                graph.render(self.topo_img, format='svg', cleanup=True, quiet=True)

                content = ""
                with open(svg_file, 'r', encoding='utf-8') as f:
                    content = f.read()

                pattern = r'<!--\s*Generated by graphviz.*?-->'
                mod_content = re.sub(pattern, '', content, flags=re.DOTALL)
                with open(svg_file, 'w', encoding='utf-8') as f:
                    f.write(mod_content)
            except graphviz.backend.execute.CalledProcessError as e:
                msg = e.stderr.decode()
                print(f"Failed rendering SVG from {self.topo_dot}: {msg}")
            except UnboundLocalError:
                print(f"Failed cleaning {svg_file}, empty or missing.")

    def __replace_in(self, steps, variables):
        result = []
        for step in steps:
            new_step = step
            for k, v in variables.items():
                new_step = new_step.replace("{" + k + "}", v)
            result.append(new_step)
        return result

    def gen_spec(self, title, variables):
        """Generate a AsciiDoc specification for the test case"""
        with open(self.test_path, 'r', encoding='utf-8') as file:
            script_content = file.read()

        parsed_script = ast.parse(script_content)
        visitor = TestStepVisitor()
        visitor.visit(parsed_script)
        description = visitor.description if visitor.description != "" else "Undefined"
        test_steps = visitor.test_steps

        if variables is not None:
            for k, v in variables.items():
                if "{" + k + "}" in description:
                    description = description.replace("{" + k + "}", v)

                test_steps = self.__replace_in(test_steps, variables)

        if title is None:
            title = visitor.name

        self.gen_topology()

        with open(self.spec_path, "w", encoding='utf-8') as spec:
            # This is the test name/title  for the test-specification.pdf,
            # it is replaced by 9PM with the 'name:' from the suite.yaml
            # when it generates the test report.
            spec.write(f"=== {title}\n\n")

            # When included in another document, e.g., the test-report.pdf,
            # AsciiDoc needs to know where iamges are located.  For other
            # use-cases, e.g., browsing in GitHub '.' is implied.
            spec.write(f"ifdef::topdoc[:imagesdir: {{topdoc}}../../{self.test_dir}]\n\n")

            spec.write("==== Description\n\n")
            spec.write(description + "\n\n")

            spec.write("==== Topology\n\n")
            spec.write(f"image::topology.svg[{title} topology, align=center, scaledwidth=75%]\n\n")

            spec.write("==== Sequence\n\n")
            spec.writelines([f". {step}\n" for step in test_steps])
            spec.write("\n\n")


def parse_suite(directory, root, suitefile):
    """Parse 9pm .yaml suite file"""
    readme = None
    first_test = True

    log(f"Opening suite file {suitefile} ...")
    with open(suitefile, "r", encoding='utf-8') as f:
        data = yaml.safe_load(f)

    for entry in data:
        if entry.get("suite"):
            path = os.path.dirname(entry["suite"])
            parse_suite(f"{directory}/{path}", root, f"{directory}/{entry['suite']}")
        elif entry.get("case"):
            name = os.path.basename(entry["case"])
            base = os.path.splitext(name)[0]
            path = os.path.dirname(entry["case"])
            opts = None

            log(f"Parsing test case {name} in {path} ...")
            if entry.get("infamy"):
                specification = entry["infamy"].get("specification", True)
                if not specification:
                    continue

            if entry.get("name"):
                test_title = entry["name"]
            else:
                test_title = None  # Set from the docstring in the test case

            if readme is None:
                readme = open(f"{directory}/{path}/Readme.adoc.tmp", "w", encoding='utf-8')

            if "opts" in entry:
                opts = dict((k[2:], v) for k, v in zip(entry["opts"][::2], entry["opts"][1::2]))

            test_case = TestCase(f"{directory}/{path}", name, root)
            test_case.gen_spec(test_title, opts)

            if first_test:
                first_test = False
            else:
                readme.write("<<<\n\n")  # Page break before subsequent tests
            readme.write(f"include::{base}.adoc[]\n\n")

            if path != "":
                lnk = f"{directory}/{path}/Readme.adoc"
                readme.close()
                log(f"Removing {directory}/{path}/Readme.adoc.tmp")
                os.unlink(f"{directory}/{path}/Readme.adoc.tmp")
                readme = None
                log(f"Checking if lnk {lnk} exists ...")
                if os.path.lexists(lnk):
                    log(f"Removing {lnk}")
                    os.unlink(lnk)
                log(f"Creating symlink {base}.adoc -> {lnk}")
                os.symlink(f"{base}.adoc", lnk)

    if readme is not None:
        readme.close()
        os.rename(f"{directory}/Readme.adoc.tmp", f"{directory}/Readme.adoc")


parser = argparse.ArgumentParser(description="Generate a test specification for a subtree.")
parser.add_argument("-s", "--suite", required=True, help="The suite file to parse")
parser.add_argument("-r", "--root-dir", required=True, help="Path that all paths should be relative to")
parser.add_argument("-d", "--debug", help="Show debug output", action="store_true")
args = parser.parse_args()

# This is hacky, graphviz output error only to stdout and *always*
# return successful.  If everything goes well, output shall be empty,
# fail on any output.
ERROR_STRING = ""

if not args.debug:
    output_capture = io.StringIO()
    sys.stderr = output_capture
    output_capture.truncate(0)
    output_capture.seek(0)

parse_suite(os.path.dirname(args.suite), args.root_dir, args.suite)

if not args.debug:
    sys.stdout = sys.__stdout__

    if len(ERROR_STRING) > 0:
        print(ERROR_STRING)
        sys.exit(1)

sys.exit(0)