"""CSV export for spreadsheet analysis.

Exports trace data and metrics to CSV format for:
- Analysis in spreadsheet applications
- Data processing with pandas/numpy
- Reporting and visualization tools
"""

import csv
from io import StringIO
from pathlib import Path
from typing import Optional

from wafer_core.lib.distributed_traces.models.trace_session import (
    AnalysisResult,
    CollectiveMatch,
    StragglerReport,
    TraceSession,
)


def export_collectives_csv(
    session: TraceSession,
    output_path: Optional[str | Path] = None,
) -> str:
    """Export all collectives to CSV.

    Args:
        session: TraceSession to export
        output_path: Optional path to write to

    Returns:
        CSV string
    """
    output = StringIO()
    writer = csv.writer(output)

    # Write header
    writer.writerow([
        "rank",
        "collective_type",
        "start_time_ns",
        "end_time_ns",
        "duration_ns",
        "duration_ms",
        "message_size_bytes",
        "count",
        "datatype",
        "process_group",
        "sequence_number",
        "algorithm",
        "protocol",
        "pytorch_op_name",
    ])

    # Write data
    for timeline in session.timelines.values():
        for c in timeline.collectives:
            writer.writerow([
                c.rank,
                c.collective_type.value,
                c.start_time_ns,
                c.end_time_ns,
                c.duration_ns,
                f"{c.duration_ms:.3f}",
                c.message_size_bytes,
                c.count,
                c.datatype.value,
                c.process_group.comm_id if c.process_group else "",
                c.sequence_number or "",
                c.algorithm or "",
                c.protocol or "",
                c.pytorch_op_name or "",
            ])

    csv_str = output.getvalue()

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

    return csv_str


def export_metrics_csv(
    result: AnalysisResult,
    matches: list[CollectiveMatch],
    output_path: Optional[str | Path] = None,
) -> str:
    """Export per-collective metrics to CSV.

    Args:
        result: AnalysisResult from analysis
        matches: Collective matches across ranks
        output_path: Optional path to write to

    Returns:
        CSV string
    """
    output = StringIO()
    writer = csv.writer(output)

    # Write header
    writer.writerow([
        "sequence",
        "collective_type",
        "process_group",
        "message_size_bytes",
        "num_ranks",
        "min_duration_ms",
        "max_duration_ms",
        "median_duration_ms",
        "straggler_rank",
        "straggler_impact_ms",
    ])

    # Write data
    for match in matches:
        writer.writerow([
            match.sequence_number,
            match.collective_type.value,
            match.process_group_id,
            match.message_size_bytes,
            match.num_ranks,
            f"{match.min_duration_ns / 1_000_000:.3f}",
            f"{match.max_duration_ns / 1_000_000:.3f}",
            f"{match.median_duration_ns / 1_000_000:.3f}",
            match.straggler_rank,
            f"{match.straggler_impact_ns / 1_000_000:.3f}",
        ])

    csv_str = output.getvalue()

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

    return csv_str


def export_straggler_csv(
    report: StragglerReport,
    output_path: Optional[str | Path] = None,
) -> str:
    """Export straggler report to CSV.

    Args:
        report: StragglerReport from analysis
        output_path: Optional path to write to

    Returns:
        CSV string
    """
    output = StringIO()
    writer = csv.writer(output)

    # Write per-rank scores
    writer.writerow(["Rank Straggler Scores"])
    writer.writerow(["rank", "score", "is_flagged"])

    for rank, score in sorted(report.rank_scores.items()):
        is_flagged = rank in report.flagged_ranks
        writer.writerow([rank, f"{score:.0f}", is_flagged])

    writer.writerow([])

    # Write top impactful collectives
    writer.writerow(["Top Impactful Collectives"])
    writer.writerow([
        "sequence",
        "collective_type",
        "straggler_rank",
        "impact_ns",
        "impact_ms",
    ])

    for match, rank in report.top_collectives[:20]:
        impact = match.straggler_impact_ns
        writer.writerow([
            match.sequence_number,
            match.collective_type.value,
            rank,
            impact,
            f"{impact / 1_000_000:.3f}",
        ])

    writer.writerow([])

    # Write summary
    writer.writerow(["Summary"])
    writer.writerow(["Total straggler impact (ms)", f"{report.total_straggler_impact_ns / 1_000_000:.3f}"])
    writer.writerow(["Flagged ranks", ", ".join(map(str, report.flagged_ranks))])
    writer.writerow(["Is consistent straggler", report.is_consistent])

    csv_str = output.getvalue()

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

    return csv_str


def export_overlap_csv(
    result: AnalysisResult,
    output_path: Optional[str | Path] = None,
) -> str:
    """Export overlap analysis to CSV.

    Args:
        result: AnalysisResult containing overlap metrics
        output_path: Optional path to write to

    Returns:
        CSV string
    """
    output = StringIO()
    writer = csv.writer(output)

    metrics = result.overlap_metrics

    writer.writerow(["Overlap Analysis"])
    writer.writerow([])
    writer.writerow(["Metric", "Value (ns)", "Value (ms)", "Percentage"])

    total_time = (
        metrics.compute_only_time_ns
        + metrics.communication_only_time_ns
        + metrics.overlapped_time_ns
        + metrics.idle_time_ns
    )

    def pct(val: int) -> str:
        return f"{(val / total_time * 100):.1f}%" if total_time > 0 else "0%"

    writer.writerow([
        "Total Compute Time",
        metrics.total_compute_time_ns,
        f"{metrics.total_compute_time_ns / 1_000_000:.3f}",
        "",
    ])
    writer.writerow([
        "Total Communication Time",
        metrics.total_communication_time_ns,
        f"{metrics.total_communication_time_ns / 1_000_000:.3f}",
        "",
    ])
    writer.writerow([
        "Overlapped Time",
        metrics.overlapped_time_ns,
        f"{metrics.overlapped_time_ns / 1_000_000:.3f}",
        pct(metrics.overlapped_time_ns),
    ])
    writer.writerow([
        "Compute-Only Time",
        metrics.compute_only_time_ns,
        f"{metrics.compute_only_time_ns / 1_000_000:.3f}",
        pct(metrics.compute_only_time_ns),
    ])
    writer.writerow([
        "Communication-Only Time",
        metrics.communication_only_time_ns,
        f"{metrics.communication_only_time_ns / 1_000_000:.3f}",
        pct(metrics.communication_only_time_ns),
    ])
    writer.writerow([
        "Idle Time",
        metrics.idle_time_ns,
        f"{metrics.idle_time_ns / 1_000_000:.3f}",
        pct(metrics.idle_time_ns),
    ])
    writer.writerow([])
    writer.writerow([
        "Overlap Ratio",
        "",
        "",
        f"{metrics.overlap_percent:.1f}%",
    ])

    # Per-iteration if available
    if metrics.iteration_metrics:
        writer.writerow([])
        writer.writerow(["Per-Iteration Overlap"])
        writer.writerow([
            "iteration",
            "compute_ms",
            "communication_ms",
            "overlap_ms",
            "overlap_ratio",
        ])

        for it in metrics.iteration_metrics:
            writer.writerow([
                it.iteration,
                f"{it.compute_time_ns / 1_000_000:.3f}",
                f"{it.communication_time_ns / 1_000_000:.3f}",
                f"{it.overlapped_time_ns / 1_000_000:.3f}",
                f"{it.overlap_ratio:.2f}",
            ])

    csv_str = output.getvalue()

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

    return csv_str
