# Usage:
#   vivado -mode batch -source run.tcl -tclargs <hls_ip_repo_dir> <project_work_dir> [jobs] [out_dir]
#
# Args:
#   $1 = Vitis HLS IP repository directory, normally <staged-hls>/export/export
#   $2 = fresh Vivado project work directory
#   $3 = parallel jobs, optional, default 8
#   $4 = output directory, optional, default <project_work_dir>/out

proc fail {message} {
  puts stderr "ERROR: $message"
  exit 1
}

proc find_files_recursive {root pattern} {
  set matches {}
  foreach path [glob -nocomplain -directory $root *] {
    if {[file isdirectory $path]} {
      set matches [concat $matches [find_files_recursive $path $pattern]]
    } elseif {[string match $pattern [file tail $path]]} {
      lappend matches $path
    }
  }
  return $matches
}

proc require_run_complete {run_name} {
  set run_obj [get_runs -quiet $run_name]
  if {[llength $run_obj] == 0} {
    fail "Run '$run_name' does not exist."
  }

  set status [get_property STATUS $run_obj]
  if {![string match "*Complete*" $status]} {
    fail "Run '$run_name' did not complete successfully. Status: $status"
  }
}

if {[llength $argv] < 2} {
  fail "Expected: run.tcl <hls_ip_repo_dir> <project_work_dir> [jobs] [out_dir]"
}

set hls_ip_repo_dir [file normalize [lindex $argv 0]]
set project_work_dir [file normalize [lindex $argv 1]]
set njobs [expr {[llength $argv] >= 3 ? [lindex $argv 2] : 8}]
set requested_outdir [expr {[llength $argv] >= 4 ? [lindex $argv 3] : ""}]

set script_dir [file dirname [file normalize [info script]]]
set bd_script [file join $script_dir create_bd.tcl]

if {![file isdirectory $hls_ip_repo_dir]} {
  fail "HLS IP repository not found: $hls_ip_repo_dir"
}
if {![file isfile $bd_script]} {
  fail "Block-design Tcl not found: $bd_script"
}

if {$requested_outdir eq ""} {
  set outdir [file normalize [file join $project_work_dir out]]
} else {
  set outdir [file normalize $requested_outdir]
}

file delete -force $project_work_dir
file mkdir $project_work_dir
file mkdir $outdir

puts "Creating Vivado project in: $project_work_dir"
create_project encoder $project_work_dir -part xczu9eg-ffvb1156-2-e -force
set_property board_part xilinx.com:zcu102:part0:3.4 [current_project]
set_property target_language Verilog [current_project]

puts "Using HLS IP repository: $hls_ip_repo_dir"
set_property ip_repo_paths [list $hls_ip_repo_dir] [current_project]
set_property ip_repo_paths [list $hls_ip_repo_dir] [get_filesets sources_1]
update_ip_catalog -rebuild

puts "Creating block design from: $bd_script"
if {[catch {source $bd_script} bd_error]} {
  fail "Block-design creation failed: $bd_error"
}

set top_bd_design [current_bd_design -quiet]
if {$top_bd_design eq ""} {
  fail "Block-design script did not leave a top-level block design open."
}
set top_bd_name [get_property NAME $top_bd_design]
set project_name [get_property NAME [current_project]]
set project_dir [get_property DIRECTORY [current_project]]
set top_bd_file [file join $project_dir "${project_name}.srcs" sources_1 bd $top_bd_name "${top_bd_name}.bd"]

if {![file exists $top_bd_file]} {
  fail "Top-level block design not found: $top_bd_file"
}

set top_bd_obj [get_files -quiet $top_bd_file]
if {[llength $top_bd_obj] == 0} {
  fail "Top-level block design is not registered in the project: $top_bd_file"
}

puts "Generating top-level block-design targets for: $top_bd_file"
open_bd_design $top_bd_file
validate_bd_design
save_bd_design
generate_target all $top_bd_obj

set wrappers [make_wrapper -files $top_bd_obj -top]
foreach wrapper $wrappers {
  if {[llength [get_files -quiet $wrapper]] == 0} {
    add_files -norecurse $wrapper
  }
}

set_property top encoder_wrapper [current_fileset]
update_compile_order -fileset sources_1

puts "Launching synthesis with $njobs job(s)."
launch_runs synth_1 -jobs $njobs
wait_on_run synth_1
require_run_complete synth_1

puts "Launching implementation and bitstream generation with $njobs job(s)."
launch_runs impl_1 -to_step write_bitstream -jobs $njobs
wait_on_run impl_1
require_run_complete impl_1

set impl_dir [get_property DIRECTORY [get_runs impl_1]]
set topname [get_property top [current_fileset]]
set primary_bit [file join $impl_dir "${topname}.bit"]
if {![file exists $primary_bit]} {
  fail "Expected primary bitstream not found: $primary_bit"
}

set produced_files {}
set bit_dst [file join $outdir encoder.bit]
file copy -force $primary_bit $bit_dst
lappend produced_files $bit_dst

set xsa_path [file join $outdir encoder.xsa]
if {[catch {write_hw_platform -fixed -include_bit -force -file $xsa_path} xsa_msg]} {
  puts stderr "WARN: Failed to export hardware platform XSA:\n$xsa_msg"
  set xsa_path ""
} elseif {[file exists $xsa_path]} {
  lappend produced_files $xsa_path
}

set top_level_hwh ""
set hwh_candidates [find_files_recursive $project_work_dir "*.hwh"]
foreach candidate $hwh_candidates {
  set tail_name [file tail $candidate]
  if {$tail_name eq "encoder.hwh" || $tail_name eq "${topname}.hwh" || $tail_name eq "${topname}_wrapper.hwh"} {
    set top_level_hwh $candidate
    break
  }
}
if {$top_level_hwh eq "" && [llength $hwh_candidates] == 1} {
  set top_level_hwh [lindex $hwh_candidates 0]
}

if {$top_level_hwh ne "" && [file exists $top_level_hwh]} {
  set hwh_dst [file join $outdir encoder.hwh]
  file copy -force $top_level_hwh $hwh_dst
  lappend produced_files $hwh_dst
} else {
  puts "WARN: Could not locate a standalone HWH artifact. FAT will try encoder.xsa as a fallback."
}

puts "\nArtifacts copied to: $outdir"
foreach f $produced_files { puts "  - $f" }

close_project
puts "Done."
puts "Primary bitstream: $primary_bit"
puts "Primary HWH: $top_level_hwh"
puts "Primary XSA: $xsa_path"
