"""JSON export for pipeline integration.

Exports trace sessions and analysis results to JSON format for:
- CI/CD pipeline integration
- Data exchange with other tools
- Archival and reproducibility
"""

import json
from dataclasses import asdict
from pathlib import Path
from typing import Any, Optional

from wafer_core.lib.distributed_traces.models.collective import CollectiveType, DataType
from wafer_core.lib.distributed_traces.models.trace_session import (
    AnalysisResult,
    HangRiskLevel,
    TraceFormat,
    TraceSession,
)


def export_session_json(
    session: TraceSession,
    output_path: Optional[str | Path] = None,
    include_events: bool = True,
) -> str:
    """Export a trace session to JSON.

    Args:
        session: TraceSession to export
        output_path: Optional path to write to
        include_events: If True, include all collectives and compute events

    Returns:
        JSON string
    """
    data = _session_to_dict(session, include_events)
    json_str = json.dumps(data, indent=2, default=_json_default)

    if output_path:
        Path(output_path).write_text(json_str)

    return json_str


def export_analysis_json(
    result: AnalysisResult,
    session: TraceSession,
    output_path: Optional[str | Path] = None,
    include_events: bool = True,
) -> str:
    """Export analysis results to JSON.

    Args:
        result: AnalysisResult from analysis
        session: TraceSession that was analyzed
        output_path: Optional path to write to
        include_events: If True, include all collectives and compute events in timelines

    Returns:
        JSON string
    """
    data = {
        "session": _session_to_dict(session, include_events=include_events),
        "analysis": _analysis_to_dict(result),
    }

    json_str = json.dumps(data, indent=2, default=_json_default)

    if output_path:
        Path(output_path).write_text(json_str)

    return json_str


def _session_to_dict(session: TraceSession, include_events: bool) -> dict:
    """Convert a TraceSession to a dictionary."""
    data = {
        "name": session.name,
        "trace_format": session.trace_format.value,
        "world_size": session.world_size,
        "num_ranks": session.num_ranks,
        "ranks": session.ranks,
        "total_collectives": session.total_collectives,
        "duration_ns": session.duration_ns,
        "duration_ms": session.duration_ns / 1_000_000,
        "aligned": session.aligned,
        "metadata": session.metadata,
    }

    if include_events:
        data["timelines"] = {
            rank: _timeline_to_dict(timeline)
            for rank, timeline in session.timelines.items()
        }

    return data


def _timeline_to_dict(timeline: "RankTimeline") -> dict:  # noqa: F821
    """Convert a RankTimeline to a dictionary."""
    from wafer_core.lib.distributed_traces.models.rank_timeline import RankTimeline

    return {
        "rank": timeline.rank,
        "device_id": timeline.device_id,
        "hostname": timeline.hostname,
        "pid": timeline.pid,
        "num_collectives": len(timeline.collectives),
        "num_compute_events": len(timeline.compute_events),
        "total_collective_time_ms": timeline.total_collective_time_ns / 1_000_000,
        "total_compute_time_ms": timeline.total_compute_time_ns / 1_000_000,
        "collectives": [_collective_to_dict(c) for c in timeline.collectives],
        "compute_events": [_compute_to_dict(e) for e in timeline.compute_events],
    }


def _collective_to_dict(collective: "Collective") -> dict:  # noqa: F821
    """Convert a Collective to a dictionary."""
    from wafer_core.lib.distributed_traces.models.collective import Collective

    return {
        "collective_type": collective.collective_type.value,
        "rank": collective.rank,
        "start_time_ns": collective.start_time_ns,
        "end_time_ns": collective.end_time_ns,
        "duration_ns": collective.duration_ns,
        "duration_ms": collective.duration_ms,
        "message_size_bytes": collective.message_size_bytes,
        "count": collective.count,
        "datatype": collective.datatype.value,
        "process_group": (
            {
                "comm_id": collective.process_group.comm_id,
                "name": collective.process_group.name,
                "world_size": collective.process_group.world_size,
            }
            if collective.process_group
            else None
        ),
        "sequence_number": collective.sequence_number,
        "algorithm": collective.algorithm,
        "protocol": collective.protocol,
        "pytorch_op_name": collective.pytorch_op_name,
    }


def _compute_to_dict(event: "ComputeEvent") -> dict:  # noqa: F821
    """Convert a ComputeEvent to a dictionary."""
    from wafer_core.lib.distributed_traces.models.rank_timeline import ComputeEvent

    return {
        "name": event.name,
        "start_time_ns": event.start_time_ns,
        "end_time_ns": event.end_time_ns,
        "duration_ns": event.duration_ns,
        "duration_ms": event.duration_ms,
        "stream_id": event.stream_id,
        "device_id": event.device_id,
        "grid_size": event.grid_size,
        "block_size": event.block_size,
    }


def _analysis_to_dict(result: AnalysisResult) -> dict:
    """Convert AnalysisResult to a dictionary."""
    return {
        "aggregate_bandwidth": {
            "achieved_gbps": result.aggregate_bandwidth.achieved_bandwidth_gbps,
            "theoretical_gbps": result.aggregate_bandwidth.theoretical_bandwidth_gbps,
            "efficiency_percent": result.aggregate_bandwidth.efficiency_percent,
            "total_bytes": result.aggregate_bandwidth.message_size_bytes,
            "total_duration_ns": result.aggregate_bandwidth.duration_ns,
        },
        "overlap_metrics": {
            "total_compute_time_ms": result.overlap_metrics.total_compute_time_ns / 1_000_000,
            "total_communication_time_ms": result.overlap_metrics.total_communication_time_ns / 1_000_000,
            "overlapped_time_ms": result.overlap_metrics.overlapped_time_ns / 1_000_000,
            "overlap_ratio": result.overlap_metrics.overlap_ratio,
            "overlap_percent": result.overlap_metrics.overlap_percent,
        },
        "straggler_report": {
            "rank_scores": result.straggler_report.rank_scores,
            "flagged_ranks": result.straggler_report.flagged_ranks,
            "total_impact_ms": result.straggler_report.total_straggler_impact_ns / 1_000_000,
            "is_consistent": result.straggler_report.is_consistent,
        },
        "hang_risk_report": {
            "timeout_threshold_sec": result.hang_risk_report.timeout_threshold_ns / 1_000_000_000,
            "risk_distribution": {
                level.value: count
                for level, count in result.hang_risk_report.risk_distribution.items()
            },
            "recommendations": result.hang_risk_report.recommendations,
        },
        "collective_matches": [_collective_match_to_dict(m) for m in result.collective_matches],
        "num_matched_collectives": len(result.collective_matches),
        "summary": result.summary,
    }


def _collective_match_to_dict(match: "CollectiveMatch") -> dict:  # noqa: F821
    """Convert a CollectiveMatch to a dictionary."""
    from wafer_core.lib.distributed_traces.models import CollectiveMatch

    return {
        "sequence": match.sequence_number,
        "collective_type": match.collective_type.value if hasattr(match.collective_type, 'value') else str(match.collective_type),
        "straggler_rank": match.straggler_rank,
        "process_group_id": match.process_group_id,
        "message_size_bytes": match.message_size_bytes,
        "num_ranks": match.num_ranks,
        "min_duration_ns": match.min_duration_ns,
        "max_duration_ns": match.max_duration_ns,
        "median_duration_ns": match.median_duration_ns,
        "straggler_impact_ns": match.straggler_impact_ns,
        "points": [
            {
                "rank": rank,
                "start_ns": collective.start_time_ns,
                "end_ns": collective.end_time_ns,
                "duration_ns": collective.duration_ns,
            }
            for rank, collective in sorted(match.collectives.items())
        ],
    }


def _json_default(obj: Any) -> Any:
    """Default JSON serializer for objects not serializable by default."""
    if isinstance(obj, (CollectiveType, DataType, HangRiskLevel, TraceFormat)):
        return obj.value
    if hasattr(obj, "__dict__"):
        return obj.__dict__
    return str(obj)
