__author__ = "7987847, Werner"
__email__ = "s5260822@stud.uni-frankfurt.de"

import os
import re
import sys

DEBUG_PRINTS = False

def debug_print(*args, **kwargs) -> None:
    """
    Print debug messages if the DEBUG_PRINTS flag is set to True.

    :param args: The positional arguments to print.
    :param kwargs: The keyword arguments to print.
    """
    if DEBUG_PRINTS:
        print(*args, **kwargs)

def get_java_files(root_dir:str) -> list:
    """
    Get all Java files in the root_dir and its subdirectories.

    :param root_dir: The root directory to search for Java files.
    :return: A list of relative file paths to the Java files.
    """
    java_files = []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith('.java'):
                # Create relative file path from the root_dir
                relative_path = os.path.relpath(os.path.join(dirpath, filename), root_dir)
                java_files.append(relative_path)
    print(f"Found {len(java_files)} Java files in {root_dir} and its subdirectories.")
    return java_files


def get_package_slug(java_file:str) -> str:
    """
    Get the base package slug from a Java file.

    :param java_file: The relative file path to the Java file.
    :return: The base package slug.
    """
    # open the file and read the first line
    with open(java_file, "r") as file:
        line = file.readline()
        # extract the package name from the first line
        if line.startswith("package "):
            return line.split(" ")[1].replace(";", "").strip()
    return ""


def is_line_field_declaration(line:str) -> bool:
    """
    Check if a line is a field declaration.

    :param line: The line to check.
    :return: True if the line is a field declaration, False otherwise.
    """
    if "public" not in line and "private" not in line and "protected" not in line:
        return
    field_pattern = re.compile(r'^\s*(public|private|protected|\s)*\s*(static|final|\s)*\s*(\w[\w\d,<>\s]*)\s+(\w[\w\d]*)\s*(=\s*[^;]*)?;')
    match = field_pattern.match(line.strip())
    return match is not None


def get_field_visibility(line):
    """
    Checks if the line is a field declaration and returns its visibility (+, -, or #).
    """
    field_pattern = re.compile(r'^\s*(public|private|protected|\s)*\s*(static|final|\s)*\s*(\w[\w\d,<>\s]*)\s+(\w[\w\d]*)\s*(=\s*[^;]*)?;')
    match = field_pattern.match(line.strip())
    if match:
        # Check for the access modifier in the line and return the corresponding symbol
        split_line = line.split(" ")

        debug_print("FOUND FIELD WITH ACCESS MODIFIER:", split_line[0])

        if "public" in split_line:
            return "+"
        elif "private" in split_line:
            return "-"
        elif "protected" in split_line:
            return "#"
        else:
            # If there's no access modifier, assume package-private (default) -> return '-'
            return "-"
    return None  # Return None if it's not a field declaration


def is_line_method_declaration(line:str, class_type:str) -> bool:
    """
    Check if a line is a method declaration.

    :param line: The line to check.
    :return: True if the line is a method declaration, False otherwise.
    """
    # Sometimes these slip in as method declarations
    if "else if" in line:
        return False
    method_pattern = re.compile(r'^\s*(public|private|protected|\s)*\s*(static|final|\s)*\s*(\w[\w\d]*)\s+(\w[\w\d]*)\s*\(.*\)\s*(?:throws\s+\w[\w\d]*)?\s*{')

    if class_type == "interface":
        method_pattern = re.compile(r'^\s*(public|private|protected|\s)*(\w[\w\d,<>\s]*)\s+\w\w*\((\w[\w\d,<>\s]*)\s*\w*\)\s*(?:throws\s+\w[\w\d]*)?\s*;$')

    match = method_pattern.match(line.strip())
    return match is not None


def get_method_visibility(line:str) -> str:
    """
    Checks if the line is a method declaration and returns its visibility (+, -, or #).
    """
    access_modifier = line.strip().split(" ")[0]

    debug_print("FOUND METHOD WITH ACCESS MODIFIER:", access_modifier)

    if "public" in access_modifier:
        return "+"
    elif "private" in access_modifier:
        return "-"
    elif "protected" in access_modifier:
        return "#"
    else:
        # If there's no access modifier, assume package-private (default) -> return '-'
        return "-"


class JavaClass:
    def __init__(self,
                 class_name,
                 class_package,
                 class_type,
                 puml_content,
                 uml_relations,
                 class_fields,
                 class_methods):
        self.class_name = class_name
        self.class_package = class_package
        self.class_type = class_type
        self.puml_content = puml_content
        self.uml_relations = uml_relations
        self.class_fields = class_fields
        self.class_methods = class_methods

    def set_puml_content(self, puml_content):
        self.puml_content = puml_content

    def get_class_name(self):
        return self.class_name

    def get_class_package(self):
        return self.class_package

    def get_class_type(self):
        return self.class_type

    def get_puml_content(self):
        return self.puml_content

    def get_uml_relations(self):
        return self.uml_relations

    def get_class_fields(self):
        return self.class_fields

    def get_class_methods(self):
        return self.class_methods


def gen_puml_code_from_class(java_class:JavaClass, no_pkgs:bool=False) -> str:
    """
    Generate PlantUML code from a JavaClass object.

    :param java_class: The JavaClass object.
    :param no_pkgs: Do not generate package visualisation in the PlantUML diagram.
    :return: The PlantUML code as a string.
    """
    puml_code = ""
    class_name = java_class.get_class_name()
    class_package = java_class.get_class_package()
    class_type = java_class.get_class_type()
    class_fields = java_class.get_class_fields()
    class_methods = java_class.get_class_methods()
    uml_relations = java_class.get_uml_relations()

    if not no_pkgs:
        puml_code += f"package {class_package + " {"}\n"

    puml_code += f"\t{class_type} \"{class_name + "\" as " + class_package + "." + class_name + " {"}\n"

    for field, visibility in class_fields.items():
        puml_code += f"\t  {visibility} {field}\n"

    for method, visibility in class_methods.items(): # TODO: no visibility if None
        if visibility:
            puml_code += f"\t  {visibility} {method}()\n"
        else:
            puml_code += f"\t  {method}()\n"

    puml_code += "\t}\n"

    if not no_pkgs:
        puml_code += "}\n"

    for related_class, relation in uml_relations.items():
        if not related_class == "*":
            cardinalities = ("1", "1")
            puml_code += f"{class_package + "." + class_name} {"\"" + cardinalities[1] + "\""} {relation}-- {"\"" + cardinalities[0] + "\""} {related_class}\n"


    puml_code += "\n"

    java_class.set_puml_content(puml_code)

    return java_class


def get_field_name_from_line(line:str) -> str:
    """
    Get the field name from a field declaration line.

    :param line: The field declaration line.
    :return: The field name.
    """
    return line.strip().split("=")[0].strip().split(" ")[-1].replace(";", "")


def get_method_name_from_line(line:str) -> str:
    """
    Get the method name from a method declaration line.

    :param line: The method declaration line.
    :return: The method name.
    """
    return line.strip().split("(")[0].split(" ")[-1]


def class_to_puml(filename:str, base_package_slug:str) -> JavaClass:
    imported_classes = {}
    puml_content = []
    uml_relations = {}
    class_fields = {}
    class_methods = {}
    class_name = filename.split("/")[-1].replace(".java", "")
    class_package = get_package_slug(filename)

    class_type = "class"

    reached_class = False
    class_getter = False
    class_setter = False

    next_line_getter = False
    next_line_setter = False

    # open file
    with open(filename, "r") as javafile:
        for javaline in javafile:
            javaline = javaline.strip()

            debug_print("LINE:", javaline)

            if next_line_getter:
                if is_line_field_declaration(javaline):
                    field_name = get_field_name_from_line(javaline)
                    class_methods[f"get{field_name.capitalize()}"] = get_field_visibility(javaline)
                    next_line_getter = False
                elif is_line_method_declaration(javaline):
                    method_name = javaline.split(" ")[-1].replace("()", "")
                    class_methods[method_name] = get_method_visibility(javaline)
                next_line_getter = False
            if next_line_setter:
                if is_line_field_declaration(javaline):
                    field_name = get_field_name_from_line(javaline)
                    class_methods[f"set{field_name.capitalize()}"] = get_field_visibility(javaline)
                    next_line_setter = False
                elif is_line_method_declaration(javaline):
                    method_name = get_method_name_from_line(javaline)
                    class_methods[method_name] = get_method_visibility(javaline)
                next_line_setter = False
            elif javaline.startswith("import"):
                if "*" not in javaline and base_package_slug in javaline:
                    importline_package_slug = javaline.strip().split(" ")[1].replace(";", "")
                    imported_classes[importline_package_slug] = importline_package_slug

            if javaline.startswith("@"):
                if "@Getter" in javaline and "@Setter" in javaline:
                    if not reached_class:
                        class_getter = True
                        class_setter = True
                    else:
                        next_line_getter = True
                        next_line_setter = True
                elif "@Getter" in javaline:
                    if not reached_class:
                        class_getter = True
                    else:
                        next_line_getter = True
                elif "@Setter" in javaline:
                    if not reached_class:
                        class_setter = True
                    else:
                        next_line_setter = True
            elif " enum " in javaline:
                class_type = "enum"
            elif " interface " in javaline:
                class_type = "interface"
                reached_class = True
            elif " class " in javaline:
                reached_class = True
            elif is_line_field_declaration(javaline):
                visibility = get_field_visibility(javaline)
                field_name = get_field_name_from_line(javaline)
                class_fields[field_name] = visibility

                if class_getter:
                    class_methods[f"get{field_name.capitalize()}"] = visibility
                if class_setter:
                    class_methods[f"set{field_name.capitalize()}"] = visibility

            elif is_line_method_declaration(javaline, class_type=class_type):
                visibility = get_method_visibility(javaline)
                method_name = get_method_name_from_line(javaline)
                if method_name != class_name:
                    class_methods[method_name] = visibility

    for class_package_slug in imported_classes:
        if not class_package_slug in uml_relations:
            uml_relations[class_package_slug] = "<"

    java_class = JavaClass(
        class_name,
        class_package,
        class_type,
        puml_content,
        uml_relations,
        class_fields,
        class_methods)

    return gen_puml_code_from_class(java_class)


def generate_puml(root_dir:str=".", no_pkgs:bool=False) -> None:
    """
    Generate a PlantUML diagram from the Java files in the root_dir.

    :param root_dir: The root directory to search for Java files.
    :param no_pkgs: Do not generate package visualisation in the PlantUML diagram
    """
    java_files = get_java_files(root_dir)
    if len(java_files) == 0:
        print("Error: No Java Files Found in this or any subsequent directory!")

    base_package_slug = ""

    if any("Main.java" in s for s in java_files):
        base_package_slug = get_package_slug(
            next((s for s in java_files if "Main.java" in s), None))
    else:
        base_package_slug = get_package_slug(java_files[0])

    puml_code = ""

    for class_file in java_files:
        puml_code += class_to_puml(class_file, base_package_slug).get_puml_content()

    # write the PlantUML code to a file
    with open(f"{root_dir}/generated_class_diagram.puml", "w") as puml_file:
        puml_file.write(f"@startuml \n title {base_package_slug}\n")
        puml_file.write(puml_code)
        puml_file.write("@enduml\n")


    print(f"PlantUML written to {root_dir}/generated_class_diagram.puml")



def get_arguments(argv: list) -> tuple:
    """
    Get the root directory and optional flags from the command line arguments.

    :param argv: The command line arguments.
    :return: A tuple containing the root directory and the no_pkgs flag.
    """
    root_dir = "."
    if "--dir" in argv:
        root_dir = argv[argv.index("--dir") + 1]
    no_pkgs = "--no-pkgs" in argv

    global DEBUG_PRINTS
    DEBUG_PRINTS = "--debug" in argv
    return root_dir, no_pkgs


def print_help() -> None:
    """
    Print a help message.
    """
    print("Usage: python puml_generator.py")
    print("Generates a PlantUML diagram from the Java files in the root_dir.")
    print("Optional flags:")
    print("--dir <root_dir> \t\t Specify the root directory for the Java files.")
    print("--help \t\t\t\t Print this help message.")
    print("--no-pkgs \t\t\t Do not generate package visualisation in the PlantUML diagram.")

if __name__ == "__main__":
    if "-h" in sys.argv or "--help" in sys.argv:
        print_help()
    else:
        root_dir, no_pkgs = get_arguments(sys.argv)
        generate_puml(root_dir, no_pkgs)