import math
from datetime import datetime as dt

from oteldatareplay.common import gen_span_id, gen_trace_id


def dt2span_time_nano(t):
    return str(math.floor(t.timestamp() * 1e9))


def span_start_time(span, key='startTimeUnixNano'):
    return dt.fromtimestamp(int(span[key]) / 1e9)


class ResourceSpanRecord:
    def __init__(self, record):
        self.record = record
        self.parse_stats()
        self.traces_ids = []
        self.span_ids = []
        self.span_count = 0
        self.first_span_start_time = None
        self.parse_stats()

    def get_first_span(self):
        rs = self.get_first_resource_span()
        if not rs:
            return None
        scope_spans = rs['scopeSpans']
        scope_span = scope_spans[0]
        spans = scope_span['spans']
        return spans[0] if len(spans) > 0 else None

    def get_first_resource_span(self):
        rss = self.get_resource_spans()
        return rss[0] if len(rss) > 0 else None

    def get_resource_spans(self):
        return self.record['resourceSpans']

    def parse_stats(self):
        rss = self.get_resource_spans()
        scope_spans = [sc for rs in rss for sc in rs['scopeSpans']]
        spans = [span for scopeSpan in scope_spans for span in scopeSpan['spans']]
        self.traces_ids = set([s['traceId'] for s in spans])
        self.span_ids = set([s['spanId'] for s in spans])
        self.span_count = len(spans)
        self.first_span_start_time = dt.fromtimestamp(int(spans[0]['startTimeUnixNano']) / 1e9)

    def iter_spans(self, func):
        for rs in self.record['resourceSpans']:
            for sc in rs['scopeSpans']:
                for span in sc['spans']:
                    func(span)

    def update_span(self, now, first_record_time, trace_id_map):
        def update_span_time(span):
            span_start_dt = span_start_time(span)
            span_end_dt = span_start_time(span, 'endTimeUnixNano')
            span_start_dt_sim = (span_start_dt - first_record_time) + now
            span_end_dt_sim = (span_end_dt - first_record_time) + now
            trace_id = span['traceId']
            if trace_id in trace_id_map:
                span['traceId'] = trace_id_map[trace_id]
            else:
                tid = gen_trace_id()
                trace_id_map[trace_id] = gen_trace_id()
                span['traceId'] = tid
            span['startTimeUnixNano'] = dt2span_time_nano(span_start_dt_sim)
            span['endTimeUnixNano'] = dt2span_time_nano(span_end_dt_sim)

        self.iter_spans(update_span_time)