#!/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 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 = {
               'type':               protoType,
               'measurementId':      int(fields[1]),
               'sourceAddress':      fields[2],
               'destinationAddress': fields[3],
               'timestamp':          formatTimestamp(fields[4], timestampAsString),
               'round':              int(fields[5]),
               'totalHops':          int(fields[6]),
               'trafficClass':       int(fields[7], 16),
               'packetSize':         int(fields[8]),
               'checksum':           int(fields[9], 16),
               'sourcePort':         int(fields[10]),
               'destinationPort':    int(fields[11]),
               'statusFlags':        int(fields[12]),
               'pathHash':           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] = {
               'sendTimestamp':   formatTimestamp(hopFields[0], timestampAsString),
               'hopNumber':       int(hopFields[1]),
               'responseSize':    int(hopFields[2]),
               'status':          int(hopFields[3]),
               'timeSource':      int(hopFields[4], 16),
               'delayAppSend':    int(hopFields[5]),
               'delayQueuing':    int(hopFields[6]),
               'delayAppReceive': int(hopFields[7]),
               'rttApp':          int(hopFields[8]),
               'rttSw':           int(hopFields[9]),
               'rttHw':           int(hopFields[10]),
               'hopAddress':      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))
