#! /usr/bin/env python
import csv
import glob
import os
from datetime import datetime
import sys

def merge_log_files(nav_file, usbl_file, output_file):
    """
    Merges navigation and USBL log files into a single CSV file with the format:
    timestamp,data,value

    Args:
        nav_file (str): Path to the navigation log CSV file.
        usbl_file (str): Path to the USBL log CSV file.
        output_file (str): Path to the output merged CSV file.
    """
    try:
        all_lines = []

        # Read the navigation log
        with open(nav_file, 'r') as f:
            reader = csv.reader(f)
            header = next(reader)  # Skip header
            for row in reader:
                all_lines.append(row)
        # Remove last line if broken
        if len(all_lines[-1]) != 3:
            all_lines.pop()

        # Read the USBL log
        with open(usbl_file, 'r') as f:
            reader = csv.reader(f)
            for row in reader:
                timestamp, data, value = row
                if timestamp == "timestamp":
                    # Remove timestamp header if present (in merge from other merge)
                    continue
                data = 'usbl_' + data
                if data == 'usbl_ID':
                    value = str(int(float(value)))
                all_lines.append([timestamp, data, value])
        if len(all_lines[-1]) != 3:
            all_lines.pop()

        # Sort by timestamp
        all_lines.sort(key=lambda x: datetime.strptime(x[0], '%Y-%m-%d %H:%M:%S.%f'))

        # Write to CSV
        with open(output_file, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['timestamp', 'data', 'value'])
            writer.writerows(all_lines)
            
        print(f"Successfully merged {nav_file} and {usbl_file} into {output_file}")

    except Exception as e:
        print(f"An error occurred: {e}")

if __name__ == "__main__":
    try:
        if len(sys.argv) == 2:
            input_file_path = sys.argv[1]
            target_directory = os.path.dirname(input_file_path)
            if target_directory == "":
                target_directory = "."

            filename_only = os.path.basename(input_file_path)
            base_name = "_".join(filename_only.split('_')[:3])
            if not base_name:
                 raise ValueError("Could not parse base name from filename.")
            print(f"Processing specific files for base name: {base_name}")

            nav_log_file = os.path.join(target_directory, f'{base_name}_navigation_log.csv')
            usbl_csv_file = os.path.join(target_directory, f'{base_name}_usbl_csv.csv')
            output_filename = os.path.join(target_directory, f'{base_name}_merged_log.csv')

            print(f"Merging {nav_log_file} and {usbl_csv_file} to {output_filename}")
            merge_log_files(nav_log_file, usbl_csv_file, output_filename)

        elif len(sys.argv) == 3:
            nav_file = sys.argv[1]
            usbl_file = sys.argv[2]

            nav_basename = os.path.splitext(os.path.basename(nav_file))[0]
            usbl_basename = os.path.splitext(os.path.basename(usbl_file))[0]
            
            output_dir = os.path.dirname(nav_file)
            if not output_dir:
                output_dir = "."
            
            output_filename = os.path.join(output_dir, f"{nav_basename}_{usbl_basename}_merged.csv")

            print(f"Merging {nav_file} and {usbl_file} to {output_filename}")
            merge_log_files(nav_file, usbl_file, output_filename)

        else:
            print("Error: Invalid number of arguments. Please provide 1 or 2 arguments.")
            print("")
            print("Usage:")
            print("  1 argument: python log_merge_nav_usbl.py <any_log_file_from_session>")
            print("     (This will find navigation and usbl logs from the same session and merge them)")
            print("  2 arguments: python log_merge_nav_usbl.py <navigation_log_file.csv> <usbl_log_file.csv>")
            print("     (This will merge the two specified files)")
            sys.exit(1)

    except FileNotFoundError as e:
        print(f"Error: File not found: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")