from copy import deepcopy
from logging import FileHandler
import os
import sys
from uuid import UUID
from typing import Dict, Union

from relationalai.debugging import JsonFormatter
from relationalai.util import get_timestamp
from relationalai.util.format import str_to_bool
from relationalai.util.span_tracker import SpanTracker, record_is_span

# Span types that are automatically generated by PyRel but are not relevant for performance analysis
SKIP_SPAN_TYPES = ['check_status']


class SpansFileHandler(FileHandler):
    """
    This logger writes spans to a local .jsonl file inside the `spans_output` folder as they finish.
    The primary goal is to ingest these spans into Snowflake for analytics.
    """

    def __init__(self, trace_id: UUID, filename: Union[str,None] = None):
        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=SKIP_SPAN_TYPES, log_span_attrs_as_str=False)
        if filename is None:
            filename = self._get_default_log_file_path()
            if not os.path.exists(os.path.dirname(filename)):
                os.makedirs(os.path.dirname(filename))
        FileHandler.__init__(self, filename=filename)
        self.setFormatter(JsonFormatter())

    def emit(self, record):
        if record_is_span(record):
            msg: Dict = record.msg # type: ignore
            if msg["event"] == "span_start":
                self.span_tracker.handle_span_start(msg)
            if msg["event"] == "span_end":
                span_end_result = self.span_tracker.handle_span_end(msg)
                if span_end_result:
                    formatted_span_msg = span_end_result[0]
                    if formatted_span_msg is None: # This will happen for span_types_to_skip
                        return
                    record_copy = deepcopy(record)  # Create a deep copy so as to not interfere with other handlers
                    record_copy.msg = formatted_span_msg
                    FileHandler.emit(self, record_copy)

    def _get_default_log_file_path(self) -> str:
        now_str = get_timestamp().strftime("%Y%m%d_%H%M%S")
        trace_id_trunc = str(self.span_tracker.trace_id)[:8]
        spans_subfolder = os.getenv("SPANS_SUBFOLDER", "")
        return os.path.join(
            os.getcwd(), "output_spans", spans_subfolder, f"{now_str}__{trace_id_trunc}.jsonl"
        )


def is_collect_only():
    return "pytest" in sys.modules and "--collect-only" in sys.argv


def maybe_add_spans_file_handler(logger, trace_id:UUID, if_not_set: bool = False):
    """
    Add a SpansFileHandler to the logger if the 'ENABLE_SPANS_FILE_HANDLER' environment variable is set to True,
    or if the 'if_not_set' parameter is True.
    """
    if str_to_bool(os.getenv("ENABLE_SPANS_FILE_HANDLER")) or if_not_set:
        if is_collect_only():
            print("Running in collect-only mode, skipping SpansFileHandler")
        elif any(isinstance(handler, SpansFileHandler) for handler in logger.handlers):
            print("logger already has a SpansFileHandler instance, won't add another one")
        else:
            logger.debug("Adding SpansFileHandler to logger")
            logger.addHandler(SpansFileHandler(trace_id))
