"""Custom OpenTelemetry SpanExporter that writes one JSON file per trace."""

import json
import os
from collections import defaultdict
from datetime import datetime, timezone
from typing import Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult


class FilePerTraceExporter(SpanExporter):
    """Exports spans grouped by trace_id into individual JSON files in a directory."""

    def __init__(self, out_dir: str):
        self.out_dir = out_dir
        os.makedirs(self.out_dir, exist_ok=True)

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        by_trace: dict[str, list[ReadableSpan]] = defaultdict(list)
        for span in spans:
            trace_id = format(span.context.trace_id, "032x")
            by_trace[trace_id].append(span)

        for trace_id, trace_spans in by_trace.items():
            existing_file = None
            for fname in os.listdir(self.out_dir):
                if fname.endswith(f"_{trace_id}.json"):
                    existing_file = os.path.join(self.out_dir, fname)
                    break

            if existing_file:
                with open(existing_file) as f:
                    data = json.load(f)
                for span in trace_spans:
                    data["spans"].append(json.loads(span.to_json()))
                with open(existing_file, "w") as f:
                    json.dump(data, f, indent=2)
            else:
                timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
                filename = f"{timestamp}_{trace_id}.json"
                filepath = os.path.join(self.out_dir, filename)
                data = {
                    "trace_id": trace_id,
                    "spans": [json.loads(span.to_json()) for span in trace_spans],
                }
                with open(filepath, "w") as f:
                    json.dump(data, f, indent=2)

        return SpanExportResult.SUCCESS

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return True
