#/*
# * Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
# *
# * Author:
# *       Nava kishore Manne <nava.kishore.manne@amd.com>
# *
# * SPDX-License-Identifier: BSD-3-Clause
# */

import struct
import sys
import types
import os
import getopt
import re
import subprocess
from pathlib import Path
from pathlib import PurePath
from lopper import Lopper
from lopper import LopperFmt
import lopper
from re import *
import yaml
import glob
from collections import OrderedDict

sys.path.append(os.path.dirname(__file__))
from baremetalconfig_xlnx import *

def is_compat( node, compat_string_to_test ):
    if re.search( "module,xlnx_overlay_dt", compat_string_to_test):
        return xlnx_generate_overlay_dt
    return ""

def infer_platform_from_sdt(sdt):
    """
    Infer the platform from the system device tree by parsing the family property.

    Args:
        sdt: System device tree object

    Returns:
        str: Inferred processor name or None if unable to determine
    """
    try:
        # Map family to processor names
        family_to_processor = {
            'Zynq': 'cortexa9',
            'ZynqMP': 'cortexa53',
            'Versal': 'cortexa72',
            'VersalNet': 'cortexa78',
            'Versal_2VE_2VM': 'cortexa78'
        }

        # Try to infer from the 'family' property at the root node
        family = sdt.tree['/'].propval('family')[0]
        if family in family_to_processor:
            processor = family_to_processor[family]
            print(f"Automatically inferred processor '{processor}' from family property '{family}'")
            return processor

        # Fallback: infer from CPU compatible strings
        for cpu_node in sdt.tree.nodes('/cpu.*@.*'):
            compatible = cpu_node.propval('compatible')[0]
            if 'cortex-a9' in compatible:
                print("Automatically inferred processor 'cortexa9' from CPU compatible string")
                return 'cortexa9'
            elif 'cortex-a53' in compatible:
                print("Automatically inferred processor 'cortexa53' from CPU compatible string")
                return 'cortexa53'
            elif 'cortex-a72' in compatible:
                print("Automatically inferred processor 'cortexa72' from CPU compatible string")
                return 'cortexa72'
            elif 'cortex-a78' in compatible:
                print("Automatically inferred processor 'cortexa78' from CPU compatible string")
                return 'cortexa78'

    except Exception as e:
        print(f"Warning: Failed to infer platform from system device tree: {e}")

    return None


def usage():
    print(f"Usage: lopper -O <output_dir> -f --enhanced <system-top.dts> <lopper-gen DT> -- xlnx_overlay_dt [<machine>] <config> <PL DTSI (SDT-gen)>")
    print("\nArguments:")
    print("  <output_dir>      - Directory where the lopper generated overlay `pl.dtsi` file will be stored")
    print("  <system-top.dts>  - Full system device tree source")
    print("  <lopper-gen DT>   - Lopper-generated system device tree file (excluding PL node)")
    print("  <machine>         - (Optional) Target architecture: cortex-a9 | cortex-a53 | cortex-a72 | cortex-a78")
    print("                      If not provided, will be automatically inferred from system device tree")
    print("  <config>          - Configuration type: full | segmented | dfx | external-fpga-config")
    print("  <PL DTSI>         - pl.dtsi file, generated by the SDT tool")

    print("\nExamples:")
    print(" With machine argument:")
    print(" lopper -O <output_dir>/ -f --enhanced <path_to_system_top>/system-top.dts <path_to_lopper_gen_dt>/lopper-gen.dts -- xlnx_overlay_dt <machine> <config> <path_to_pl_dtsi>/pl.dtsi")
    print("")
    print(" With automatic platform inference:")
    print(" lopper -O <output_dir>/ -f --enhanced <path_to_system_top>/system-top.dts <path_to_lopper_gen_dt>/lopper-gen.dts -- xlnx_overlay_dt <config> <path_to_pl_dtsi>/pl.dtsi")
    print("  --firmware-name=<name> - (Optional) Override the default firmware-name in the output")

def remove_node_ref(sdt, tgt_node, ref_node):
    prop_dict = ref_node.__props__.copy()
    match_label_list = []
    for node in sdt.tree[tgt_node].subnodes():
        matched_label = get_label(sdt, ref_node, node)
        if matched_label:
            match_label_list.append(matched_label)
    for prop,node1 in prop_dict.items():
        if prop not in match_label_list:
            sdt.tree['/' + ref_node.name].delete(prop)

"""
This API removes lines that contain `status = "okay";` from a DTS/DTSI file.

Args:
	filepath (str): Path to the input DTS/DTSI file.
"""

def remove_status_okay(filepath):
    with open(filepath, "r") as f:
        lines = f.readlines()

    modified_lines = [
        line for line in lines if 'status = "okay";' not in line.strip()
    ]

    with open(filepath, "w") as f:
        f.writelines(modified_lines)

match_list = ["v_tc", "v_smpte_uhdsdi_tx", "v_smpte_uhdsdi_rx", "v_hdmi_tx1", "v_hdmi_rx1",
	      "v_hdmi_rx", "v_hdmi_tx", "v_dp_tx", "v_dp_rx", "v_dp_rx1", "v_dp_tx1"]
"""
This API adds `status = "disabled";` after each line matching `xlnx,ip-name = "{name}";`
from the provided match_list. Preserves indentation.
Args:
	filepath (str): Path to the input DTS/DTSI file.
	match_list (List[str]): List of IP names to match.
"""
def add_status_disabled_after_ipname(filepath, match_list):
    with open(filepath, "r") as f:
        lines = f.readlines()

    modified_lines = []
    for line in lines:
        modified_lines.append(line)
        stripped_line = line.strip()

        for name in match_list:
            if stripped_line == f'xlnx,ip-name = "{name}";':
                indent = line[:len(line) - len(line.lstrip())]
                modified_lines.append(f'{indent}status = "disabled";\n')
                break

    with open(filepath, "w") as f:
        f.writelines(modified_lines)


"""
This API generates the overlay dts file by taking pl.dtsi
generated from DTG++.
Args:
    tgt_node: is the baremetal config top level domain node number
    sdt:      is the system device-tree
    options:  There are two valid invocation modes:
              Mode 1 (3 args): <machine> <config> <PL DTSI>
              Mode 2 (2 args): <config> <PL DTSI> (machine auto-inferred)

              <machine> - Target architecture: cortex-a9 | cortex-a53 | cortex-a72 | cortex-a78
              <config>  - Configuration type: full | segmented | dfx | external-fpga-config
              <PL DTSI> - pl.dtsi file, generated by the SDT tool
"""
def xlnx_generate_overlay_dt(tgt_node, sdt, options):
    try:
        # Separate positional arguments from optional arguments
        positional_args = []
        optional_args = []

        for arg in options['args']:
            if arg.startswith("--"):
                optional_args.append(arg)
            else:
                positional_args.append(arg)

        # Process optional arguments
        firmware_override = None
        for arg in optional_args:
            if arg.startswith("--firmware-name="):
                firmware_override = arg.split("=", 1)[1]

        # Determine invocation mode based on number of positional arguments
        num_args = len(positional_args)
        processor = None
        config = None
        input_file = None

        if num_args == 3:
            # Mode 1: machine, config, input_file
            processor = positional_args[0]
            config = positional_args[1].strip().lower()
            input_file = positional_args[2]

        elif num_args == 2:
            # Mode 2: config, input_file (auto-infer machine)
            config = positional_args[0].strip().lower()
            input_file = positional_args[1]

            # Auto-infer processor from system device tree
            processor = infer_platform_from_sdt(sdt)
            if processor is None:
                print("Error: Unable to automatically infer platform from system device tree.")
                print("Please provide the machine argument explicitly.")
                usage()
                return False
        else:
            print("Error: Invalid number of arguments.")
            print(f"Expected 2 or 3 positional arguments, got {num_args}")
            usage()
            return False

    except (IndexError, KeyError) as e:
        print(f"Error: Failed to parse arguments: {e}")
        usage()
        return False

    # Validate config value
    valid_configs = {"full", "segmented", "dfx", "external-fpga-config"}
    if config not in valid_configs:
        print(f"Error: Invalid config '{config}'. Expected one of: {valid_configs}.")
        return False

    print(f"Processor: {processor}, Config: {config}, Input: {input_file}")

    platform_map = {
        "cortexa9-zynq": "Zynq",
        "cortexa9_0": "Zynq",
        "cortexa9": "Zynq",
        "cortexa53-zynqmp": "ZynqMP",
        "psu_cortexa53_0": "ZynqMP",
        "cortexa53_0": "ZynqMP",
        "cortexa53": "ZynqMP",
        "cortexa72-versal": "Versal",
        "psv_cortexa72_0": "Versal",
        "cortexa72_0": "Versal",
        "cortexa72": "Versal",
        "psx_cortexa78_0": "VersalNet",
        "cortexa78_0": "VersalNet",
        "cortexa78": "VersalNet"
    }

    platform = platform_map.get(processor)
    if platform is None:
        print(f"Error: Unsupported processor '{processor}'.")
        print(f"Expected one of: {', '.join(platform_map.keys())}.")
        return False

    # Read input file
    with open(input_file, "r") as infile:
        lines = infile.readlines()

    print("Starting overlay generation...")

    # Move the amba_pl node to the overlay
    try:
        amba_node = sdt.tree["/amba_pl"]
    except KeyError:
        print("Error: /amba_pl node not found in the system device tree.")
        return False


    new_amba_node = amba_node()
    sdt.tree = sdt.tree - amba_node  # Remove from base tree

    # Remove "amba_pl" references from __symbols__ and aliases
    remove_node_ref(sdt, tgt_node, sdt.tree['/__symbols__'])
    remove_node_ref(sdt, tgt_node, sdt.tree['/aliases'])

    filtered_lines = lines[1:-1]  # Remove first and last lines

    # Lines to remove if at the top of &amba
    if platform in ["Zynq"]:
        lines_to_remove = [
            "ranges;",
            "compatible = \"simple-bus\";",
            "#address-cells = <1>;",
            "#size-cells = <1>;"
        ]
    else:
        lines_to_remove = [
            "ranges;",
            "compatible = \"simple-bus\";",
            "#address-cells = <2>;",
            "#size-cells = <2>;"
        ]

    modified_lines = []
    amba_block = []
    fpga_pr_nodes = []
    misc_clk_nodes = []
    clocking_nodes = []
    afi_nodes = []
    amba_nodes = []

    inside_fpga_pr = inside_misc_clk = inside_clocking = inside_afi = inside_amba = False
    amba_started = False
    firmware_name = None

    i = 0
    while i < len(filtered_lines):
        stripped_line = filtered_lines[i].strip()

        # Capture firmware-name
        if "firmware-name" in stripped_line:
            firmware_name = filtered_lines[i]
            i += 1
            continue

        if stripped_line.startswith("amba_pl"):
            inside_amba = True
            amba_started = True
            amba_block.append("&amba {\n")
            i += 1
            continue

        if inside_amba:
            # Remove unnecessary lines at the top
            if all(
                i + j < len(filtered_lines) and filtered_lines[i + j].strip() == lines_to_remove[j]
                for j in range(len(lines_to_remove))
            ):
                i += len(lines_to_remove)  # Skip these lines
                continue
            inside_amba = False  # Stop once a non-matching line is found

        # Handle FPGA_PR nodes
        if "fpga_PR" in stripped_line and "{" in stripped_line:
            inside_fpga_pr = True
            current_fpga_pr = [filtered_lines[i]]
            i += 1
            continue
        elif inside_fpga_pr:
            current_fpga_pr.append(filtered_lines[i])
            if "};" in stripped_line:
                fpga_pr_nodes.append(current_fpga_pr)
                inside_fpga_pr = False
            i += 1
            continue

        # Handle Misc Clock nodes
        if "misc_clk_" in stripped_line and "{" in stripped_line:
            inside_misc_clk = True
            current_misc_clk = [filtered_lines[i]]
            i += 1
            continue
        elif inside_misc_clk:
            current_misc_clk.append(filtered_lines[i])
            if "};" in stripped_line:
                misc_clk_nodes.append(current_misc_clk)
                inside_misc_clk = False
            i += 1
            continue

        # Handle Clocking nodes
        if "clocking" in stripped_line and "{" in stripped_line:
            inside_clocking = True
            current_clocking = [filtered_lines[i]]
            i += 1
            continue
        elif inside_clocking:
            current_clocking.append(filtered_lines[i])
            if "};" in stripped_line:
                clocking_nodes.append(current_clocking)
                inside_clocking = False
            i += 1
            continue

        # Handle afi nodes
        if "afi" in stripped_line and "{" in stripped_line:
            inside_afi = True
            current_afi = [filtered_lines[i]]
            i += 1
            continue
        elif inside_afi:
            current_afi.append(filtered_lines[i])
            if "};" in stripped_line:
                afi_nodes.append(current_afi)
                inside_afi = False
            i += 1
            continue

        # Store remaining amba nodes
        if amba_started:
            amba_block.append(filtered_lines[i])
        else:
            modified_lines.append(filtered_lines[i])

        i += 1

    # Append extracted FPGA_PR, misc_clk, and clocking nodes separately
    fpga_block = []
    if platform in ["ZynqMP", "Zynq"]:
        fpga_block.append("&fpga_full {\n")
    else:
        fpga_block.append("&fpga {\n")

    # Add firmware-name if present
    if platform not in ["ZynqMP", "Zynq"] and config == "dfx":
        firmware_name = "\t\texternal-fpga-config;\n"

    if config == "external-fpga-config":
        firmware_name = "\t\texternal-fpga-config;\n"

    if firmware_override:
        firmware_name = f'\t\tfirmware-name = "{firmware_override}";\n'

    if firmware_name:
        fpga_block.append(firmware_name)

    for pr_node in fpga_pr_nodes:
        fpga_block.extend(pr_node)
    for misc_node in misc_clk_nodes:
        fpga_block.extend(misc_node)
    for clocking_node in clocking_nodes:
        fpga_block.extend(clocking_node)
    for afi_node in afi_nodes:
        fpga_block.extend(afi_node)

    fpga_block.append("};\n\n")

    # Assemble final content
    final_lines = []
    final_lines.append("/dts-v1/;\n/plugin/;\n")
    final_lines.extend(fpga_block)  # FPGA-related nodes first
    final_lines.extend(modified_lines)  # General modifications
    if amba_block:
        final_lines.append("&amba {\n")  # Correct naming
        final_lines.extend(amba_block[1:])  # Avoid duplicate `&amba {`
        final_lines.append("};\n")

    # Write output file
    output_file = os.path.join(sdt.outdir, 'pl.dtsi')
    with open(output_file, "w") as outfile:
        for line in final_lines[:-1]:  # Exclude last line
            line = line[1:] if line.startswith("\t") else line  # Remove one leading tab
            line = line.replace("interrupt-parent = <&imux>;", "interrupt-parent = <&gic>;")  # Replace interrupt parent
            outfile.write(line)

    remove_status_okay(output_file)
    add_status_disabled_after_ipname(output_file, match_list)

    # Copy contents from pl.dtsi to pl.dtso in the same directory
    dtso_file = os.path.join(sdt.outdir, 'pl.dtso')
    with open(output_file, "r") as src, open(dtso_file, "w") as dst:
        dst.writelines(src.readlines())

    print("Overlay generation completed successfully!")
    return True

