#!/usr/bin/env python3
# ==========================================================================
#     _   _ _ ____            ____          _____
#    | | | (_)  _ \ ___ _ __ / ___|___  _ _|_   _| __ __ _  ___ ___ _ __
#    | |_| | | |_) / _ \ '__| |   / _ \| '_ \| || '__/ _` |/ __/ _ \ '__|
#    |  _  | |  __/  __/ |  | |__| (_) | | | | || | | (_| | (_|  __/ |
#    |_| |_|_|_|   \___|_|   \____\___/|_| |_|_||_|  \__,_|\___\___|_|
#
#       ---  High-Performance Connectivity Tracer (HiPerConTracer)  ---
#                 https://www.nntb.no/~dreibh/hipercontracer/
# ==========================================================================
#
# High-Performance Connectivity Tracer (HiPerConTracer)
# Copyright (C) 2015-2026 by Thomas Dreibholz
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact: dreibh@simula.no

import argparse
import json
import os
import sys

import HiPerConTracer

from typing import Any, Dict, Final, List, Optional, TextIO, Union


# ###### Format timestamp field #############################################
def formatTimestamp(timestampHex      : str,
                    timestampAsString : bool) -> Union[int, str]:
   unixTime : Final[int] = int(timestampHex, 16)
   if timestampAsString:
      return str(HiPerConTracer.unixtime_to_string(unixTime))
   return unixTime



# ###### Main program #######################################################

# ====== Handle arguments ===================================================
parser : Final[argparse.ArgumentParser] = argparse.ArgumentParser(
   description='Convert HiPerConTracer Traceroute results files to JSON format.'
)
parser.add_argument(
   '-T', '--timestamp-as-string',
   action = 'store_true',
   dest   = 'timestampAsString',
   help   = 'Format timestamps as strings instead of nanosecond integers'
)
parser.add_argument(
   'hpctInputFileNames',
   metavar = 'traceroute_hpct_file',
   nargs   = '+',
   help    = 'HiPerConTracer Traceroute result file(s)'
)
args               : Final[argparse.Namespace] = parser.parse_args()
timestampAsString  : Final[bool]               = args.timestampAsString
hpctInputFileNames : Final[List[str]]          = args.hpctInputFileNames

# ====== Process files ======================================================
traceroutes       : List[Dict[str, Any]]     = []
currentTraceroute : Optional[Dict[str, Any]] = None
for hpctInputFileName in hpctInputFileNames:
   # ====== Open file =======================================================
   try:
      hpctInputFile : TextIO = \
         HiPerConTracer.openHiPerConTracerFile(hpctInputFileName)
   except Exception as e:
      sys.stderr.write(f'ERROR: Unable to open input file {hpctInputFileName}: {e}\n')
      sys.exit(1)

   # ====== Parse Traceroute ================================================
   line : str
   for line in hpctInputFile:
      line = line.rstrip('\n')

      tag : str = line.split(maxsplit=1)[0]
      if tag.startswith('#T') and len(tag) == 3:
         if currentTraceroute is not None:
            traceroutes.append(currentTraceroute)
            currentTraceroute = None
         fields : List[str] = line.split()
         if len(fields) >= 14:
            protoKey  : str = fields[0][2]
            protoType : str = \
               HiPerConTracer.PROTOCOL_MAP.get(protoKey, protoKey.upper())
            currentTraceroute = {
               'service':            'Traceroute',
               'type':                protoType,
               'measurement_id':      int(fields[1]),
               'source_address':      fields[2],
               'destination_address': fields[3],
               'timestamp':           formatTimestamp(fields[4], timestampAsString),
               'round':               int(fields[5]),
               'total_hops':          int(fields[6]),
               'traffic_class':       int(fields[7], 16),
               'packet_size':         int(fields[8]),
               'checksum':            int(fields[9], 16),
               'source_port':         int(fields[10]),
               'destination_port':    int(fields[11]),
               'status_flags':        int(fields[12]),
               'path_hash':           int(fields[13], 16),
               'hops':                [ ]
            }
      elif line.startswith('\t') and currentTraceroute is not None:
         hopFields : List[str] = line.lstrip('\t').split()
         if len(hopFields) >= 12:
            currentHop : Dict[str, Any] = {
               'send_timestamp':    formatTimestamp(hopFields[0], timestampAsString),
               'hop_number':        int(hopFields[1]),
               'response_size':     int(hopFields[2]),
               'status':            int(hopFields[3]),
               'time_source':       int(hopFields[4], 16),
               'delay_app_send':    int(hopFields[5]),
               'delay_queuing':     int(hopFields[6]),
               'delay_app_receive': int(hopFields[7]),
               'rtt_app':           int(hopFields[8]),
               'rtt_sw':            int(hopFields[9]),
               'rtt_hw':            int(hopFields[10]),
               'hop_address':       hopFields[11]
            }
            currentTraceroute['hops'].append(currentHop)
      elif not tag.startswith('#?'):
         sys.stderr.write(f'ERROR: Syntax error in input file {hpctInputFileName}: {line}\n')
         sys.exit(1)

   if currentTraceroute is not None:
      traceroutes.append(currentTraceroute)
      currentTraceroute = None
   hpctInputFile.close()

# ====== Output JSON ====================================================
print(json.dumps(traceroutes, indent=3))
