#!/usr/bin/env python3
# ------------------------------------------------------------------------ #
#     (C) Copyright 2017-2025 Barcelona Supercomputing Center              #
#                             Centro Nacional de Supercomputacion          #
#                                                                          #
#     This file is part of OmpSs@FPGA toolchain.                           #
#                                                                          #
#     This code is free software; you can redistribute it and/or modify    #
#     it under the terms of the GNU Lesser General Public License as       #
#     published by the Free Software Foundation; either version 3 of       #
#     the License, or (at your option) any later version.                  #
#                                                                          #
#     OmpSs@FPGA toolchain is distributed in the hope that it will be      #
#     useful, but WITHOUT ANY WARRANTY; without even the implied           #
#     warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.     #
#     See the GNU Lesser General Public License for more details.          #
#                                                                          #
#     You should have received a copy of the GNU Lesser General Public     #
#     License along with this code. If not, see <www.gnu.org/licenses/>.   #
# ------------------------------------------------------------------------ #

import json
import os
import random
import re
import shutil
import subprocess
import sys

import ait.backend.xilinx.utils.checkers as checkers
from ait.frontend.config import BITINFO_VERSION,  \
    VERSION_MAJOR, VERSION_MINOR, VERSION_PATCH
from ait.frontend.utils import ait_path, msg, json2tcl

script_folder = os.path.basename(os.path.dirname(os.path.realpath(__file__)))


def generate_Vivado_project_tcl():
    global accs
    global args

    # Source AIT tcl procedures
    vivado_project = '# File automatically generated by the Accelerator Integration Tool. Edit at your own risk.\n' \
        + '\n' \
        + 'package require json' \
        + '\n' \
        + 'set scriptDir [file dirname [file normalize [info script]]]' \
        + '\n' \
        + 'namespace eval AIT {\n' \
        + '\tnamespace eval vars {\n' \
        + '\t\t# Project variables\n' \
        + f'\t\tset aitCall "{re.escape(os.path.basename(sys.argv[0]))} {re.escape(" ".join(sys.argv[1:]))}"\n' \
        + f'\t\tset aitMajorVersion {VERSION_MAJOR}\n' \
        + f'\t\tset aitMinorVersion {VERSION_MINOR}\n' \
        + f'\t\tset aitPatchVersion {VERSION_PATCH}\n' \
        + f'\t\tset bitinfoVersion {BITINFO_VERSION}\n'

    # Set bitstream variables
    vivado_project += '\n' \
        + '\t\t# Bitstream variables\n' \
        + f'\t\tset userID {user_id}\n'

    # Load board json file
    vivado_project += '\n' \
        + f'\t\t# Load {board.name} board info as a tcl dict\n' \
        + '\t\tset board [json::json2dict [read [open ${scriptDir}/../board/board_info.json "r"]]]\n'

    # Store accelerators information from json
    vivado_project += '\n' \
        + '\t\t# Store accelerators information in a dict\n' \
        + json2tcl(accs, 'accs', 2)

    # Load AIT configuration
    vivado_project += '\n' \
        + '\t\t# Load AIT configuration flags from json\n' \
        + '\t\tset aitConfig [json::json2dict [read [open ${scriptDir}/../../ait_config.json "r"]]]\n'

    # Load user-provided configuration
    vivado_project += '\n' \
        + '\t\t# Create userConfig dict to check for user-provided configuration\n' \
        + '\t\tif {[file exists ${scriptDir}/../user/config.json]} {\n' \
        + '\t\t\tset userConfig [json::json2dict [read [open ${scriptDir}/../user/config.json "r"]]]\n' \
        + '\t\t} else {\n' \
        + '\t\t\tset userConfig [dict create]\n' \
        + '\t\t}\n'

    vivado_project += '\t}\n'
    vivado_project += '}\n'

    vivado_project_file = open(f'{project_backend_path}/tcl/project.tcl', 'w')
    vivado_project_file.write(vivado_project)
    vivado_project_file.close()


def run_step(project_args):
    global args
    global board
    global accs
    global chip_part
    global start_time
    global user_id
    global ait_backend_path
    global project_backend_path
    global project_board_path

    args = project_args['args']
    board = project_args['board']
    accs = project_args['accs']
    start_time = project_args['start_time']
    project_path = project_args['path']

    chip_part = board.chip_part + ('-' + board.es if (board.es and not args.ignore_eng_sample) else '')
    ait_backend_path = f'{ait_path}/backend/{args.backend}'
    project_backend_path = f'{project_path}/{args.backend}'
    project_board_path = f'{project_backend_path}/board'
    user_files_path = f'{project_backend_path}/user'

    # Check if Vivado requirements are met
    checkers.check_vivado()

    # Copy AIT tcl scripts and IPs to project directory tree
    shutil.rmtree(f'{project_backend_path}/tcl', ignore_errors=True)
    shutil.copytree(f'{ait_backend_path}/tcl', f'{project_backend_path}/tcl')
    shutil.rmtree(f'{project_backend_path}/IPs', ignore_errors=True)
    shutil.copytree(f'{ait_backend_path}/IPs', f'{project_backend_path}/IPs')
    shutil.rmtree(f'{project_board_path}', ignore_errors=True)
    shutil.copytree(f'{ait_backend_path}/board/{board.name}', project_board_path)
    os.symlink(f'{project_board_path}/baseDesign.tcl', f'{project_backend_path}/tcl/templates/baseDesign.tcl')
    shutil.rmtree(f'{project_backend_path}/constraints', ignore_errors=True)
    shutil.rmtree(f'{user_files_path}', ignore_errors=True)
    os.mkdir(user_files_path)

    if args.memory_interleaving_stride > 0:
        subprocess.run(f'sed -i "s/\`undef __ENABLE__/\`define __ENABLE__/" {project_backend_path}/IPs/bsc_axiu_addrInterleaver.v', shell=True, check=True)

    if args.user_config and os.path.exists(args.user_config):
        shutil.copy2(args.user_config, f'{user_files_path}/config.json')
    elif args.user_config:
        msg.error(f'User configuration file not found: {args.user_config}')

    if args.user_constraints and os.path.exists(args.user_constraints):
        user_constraints_path = f'{user_files_path}/constraints'
        if not os.path.exists(user_constraints_path):
            os.mkdir(user_constraints_path)
        if args.verbose_info:
            msg.log(f'Adding user constraints file: {args.user_constraints}')
        shutil.copy2(args.user_constraints, user_constraints_path)
    elif args.user_constraints:
        msg.error(f'User constraints file not found: {args.user_constraints}')

    if args.user_pre_design and os.path.exists(args.user_pre_design):
        user_pre_design_ext = args.user_pre_design.split('.')[-1] if len(args.user_pre_design.split('.')) > 1 else ''
        if user_pre_design_ext != 'tcl':
            msg.error(f'Invalid extension for PRE design TCL script: {args.user_pre_design}')
        elif args.verbose_info:
            msg.log(f'Adding pre design user script: {args.user_pre_design}')
        user_tcl_scripts_path = f'{user_files_path}/tcl/scripts'
        if not os.path.exists(user_tcl_scripts_path):
            os.mkdir(user_tcl_scripts_path)
        shutil.copy2(args.user_pre_design, f'{user_files_path}/tcl/scripts/userPreDesign.tcl')
    elif args.user_pre_design:
        msg.error(f'User PRE design TCL script not found: {args.user_pre_design}')

    if args.user_post_design and os.path.exists(args.user_post_design):
        user_post_design_ext = args.user_post_design.split('.')[-1] if len(args.user_post_design.split('.')) > 1 else ''
        if user_post_design_ext != 'tcl':
            msg.error(f'Invalid extension for POST design TCL script: {args.user_post_design}')
        elif args.verbose_info:
            msg.log(f'Adding post design user script: {args.user_post_design}')
        user_tcl_scripts_path = f'{user_files_path}/tcl/scripts'
        if not os.path.exists(user_tcl_scripts_path):
            os.mkdir(user_tcl_scripts_path)
        shutil.copy2(args.user_post_design, f'{user_files_path}/tcl/scripts/userPostDesign.tcl')
    elif args.user_post_design:
        msg.error(f'User POST design TCL script not found: {args.user_post_design}')

    # Generate random USERID to identify the bitstream
    user_id = str(hex(random.randrange(2**32)))
    msg.log(f'Setting bitstream user id: {user_id}')
    subprocess.run(f'sed -i s/BITSTREAM_USERID/{user_id}/ {project_board_path}/constraints/basic_constraints.xdc', shell=True, check=True)

    # Generate tcl file with ait project variables and namespaces
    generate_Vivado_project_tcl()

    # Enable beta device on Vivado init script
    init_script_str = f'enable_beta_device {chip_part}'
    if os.path.exists(f'{project_board_path}/board_files'):
        init_script_str += f'\nset_param board.repoPaths [list {project_board_path}/board_files]'

    subprocess.run(f'echo "{init_script_str}" > {project_backend_path}/vivado.tcl', shell=True, check=True)

    p = subprocess.Popen('vivado -init -nojournal -nolog -notrace -mode batch '
                         + f'-source {project_backend_path}/tcl/project.tcl '
                         + f'-source {project_backend_path}/tcl/scripts/ait.tcl '
                         + f'-source {project_backend_path}/tcl/scripts/generate_design.tcl',
                         cwd=project_backend_path, stdout=sys.stdout.subprocess,
                         stderr=sys.stdout.subprocess, shell=True)

    if args.verbose:
        for line in iter(p.stdout.readline, b''):
            sys.stdout.write(line.decode('utf-8'))

    retval = p.wait()
    if retval:
        msg.error('Block Design generation failed', start_time, False)
    else:
        # Prettify ait json file
        with open(f'{project_path}/{args.name}.ait.json', 'r+') as ait_json_file:
            ait_json_str = json.dumps(json.load(ait_json_file), indent=4)
            ait_json_file.seek(0)
            ait_json_file.write(ait_json_str)
            ait_json_file.truncate()

        msg.success('Block Design generated')
