#!/bin/env python3
import os
import ast
import argparse
import io
import sys
import re
import yaml

from pathlib import Path

import graphviz


def replace_image_tag(text, test_dir):
    """
    Convert images added in the description and replace with the required ifdefs to work
    generating the test specifcation as well.
    """

    pattern = r"image::(?P<image_name>\S+)\[(?P<label>[^\]]*)\]"

    def repl(match):
        image_name = match.group("image_name")
        label = match.group("label")

        return f"""ifdef::topdoc[]
image::{{topdoc}}../../{test_dir}/{image_name}[{label}]
endif::topdoc[]
ifndef::topdoc[]
ifdef::testgroup[]
image::{'/'.join(Path(test_dir).parts[3:])}/{image_name}[{label}]
endif::testgroup[]
ifndef::testgroup[]
image::{image_name}[{label}]
endif::testgroup[]
endif::topdoc[]"""

    return re.sub(pattern, repl, text)

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.test_steps=[]
        self.name = ""
        self.description = ""

    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

def replace_in_steps(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

class TestCase:
    """All test specifcation resources for a test case"""
    def __init__(self, directory, rootdir=None):
        self.test_dir = Path(directory)
        if rootdir:
            rootdir = Path(f"{rootdir}")
            self.test_dir = self.test_dir.relative_to(rootdir)
        self.topology_dot = f"{directory}/topology.dot"
        self.topology_image = f"{directory}/topology"
        self.test_case = f"{directory}/test.py"
        self.specification = f"{directory}/Readme.adoc"

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

            try:
                dot_graph = dot_file.read()
                graph = graphviz.Source(dot_graph)
                graph.render(self.topology_image, 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.topology_dot}: {msg}")
            except UnboundLocalError:
                print(f"Failed cleaning {svg_file}, empty or missing.")

    def generate_specification(self, name, case_path, spec_path, variables):
        """Generate a ASCIIDOC specification for the test case"""
        with open(case_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=replace_in_steps(test_steps, variables)
        if name is None:
            name = visitor.name
        self.generate_topology()
        description = replace_image_tag(description, self.test_dir)
        with open(spec_path, "w") as spec:
            spec.write(f"=== {name}\n")
            spec.write("==== Description\n")
            spec.write(description + "\n\n")
            spec.write("==== Topology\n")
            spec.write("ifdef::topdoc[]\n")
            spec.write(f"image::{{topdoc}}../../{self.test_dir}/topology.svg[{name} topology]\n")
            spec.write("endif::topdoc[]\n")
            spec.write("ifndef::topdoc[]\n")
            spec.write("ifdef::testgroup[]\n")
            spec.write(f"image::{Path(*self.test_dir.parts[3:])}/topology.svg[{name} topology]\n")
            spec.write("endif::testgroup[]\n")
            spec.write("ifndef::testgroup[]\n")
            spec.write(f"image::topology.svg[{name} topology]\n")
            spec.write("endif::testgroup[]\n")
            spec.write("endif::topdoc[]\n")
            spec.write("==== Test sequence\n")
            spec.writelines([f". {step}\n" for step in test_steps])
            spec.write("\n\n<<<\n\n")  # need empty lines to pagebreak


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

    with open(suitefile, "r") 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"):
            variables=None
            if not entry.get("infamy"):
                test_title = None  # Set from the docstring in the test case
            else:
                specification = entry["infamy"].get("specification", True);
                if not specification:
                    continue
                test_title = entry["infamy"]["title"]
            path = os.path.dirname(entry["case"])
            if readme is None:
                readme = open(f"{directory}/{path}/Readme.adoc.tmp", "w")
            test_case_id = entry["name"]

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

            test_file = f"{directory}/{entry['case']}"
            test_spec = f"{directory}/{path}/{test_case_id}.adoc"
            test_case = TestCase(f"{directory}/{path}", root)
            test_case.generate_specification(test_title, test_file, test_spec, variables)
            readme.write(f"include::{path}{test_case_id}.adoc[]\n\n")
            if path != "":
                lnk = f"{directory}/{path}/Readme.adoc"
                readme.close()
                os.unlink(f"{directory}/{path}/Readme.adoc.tmp")
                readme = None
                if os.path.exists(lnk):
                    os.unlink(lnk)
                os.symlink(f"{test_case_id}.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()

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

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)
        exit(1)

exit(0)