"""Shared helpers for report formatting and JSON serialization.

Extracted from cli.py to reduce the monolith's size and make these
utilities reusable by other modules (badge, compare, etc.).
"""

from __future__ import annotations


def json_default(obj):
    """Handle numpy types for JSON serialization."""
    import numpy as np

    if isinstance(obj, (np.bool_,)):
        return bool(obj)
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


def coverage_status(score: float) -> tuple[str, str]:
    """Return (status, label) for a coverage score."""
    if score >= 0.65:
        return "ok", "good"
    elif score >= 0.45:
        return "warn", "moderate"
    else:
        return "bad", "needs improvement"


def score_color(score: float) -> str:
    if score >= 0.65:
        return "green"
    elif score >= 0.45:
        return "yellow"
    else:
        return "red"


def priority_style(priority: str) -> str:
    if priority == "HIGH":
        return "bold red"
    elif priority == "MEDIUM":
        return "bold yellow"
    else:
        return "dim"


def info_to_dict(info) -> dict:
    """Convert DatasetInfo to a JSON-serializable dict."""
    from orbit.utils.display import format_duration

    return {
        "repo_id": info.repo_id,
        "num_episodes": info.num_episodes,
        "num_frames": info.num_frames,
        "fps": info.fps,
        "robot_type": info.robot_type,
        "camera_names": info.camera_names,
        "video_keys": info.video_keys,
        "shapes": info.shapes,
        "duration_seconds": info.duration_seconds,
        "duration_human": format_duration(info.duration_seconds),
    }


def signal_diag_to_dict(sd) -> dict:
    """Convert a SignalDiagnosticsReport to a JSON-serializable dict."""
    return {
        "joint_diagnostics": [
            {
                "joint_index": jd.joint_index,
                "joint_name": jd.joint_name,
                "range_used": list(jd.range_used),
                "range_limit": list(jd.range_limit) if jd.range_limit else None,
                "range_utilization": round(jd.range_utilization, 4),
                "zero_velocity_ratio": round(jd.zero_velocity_ratio, 4),
                "smoothness": round(jd.smoothness, 4),
                "directional_bias": round(jd.directional_bias, 4),
                "is_dead": jd.is_dead,
                "is_clipping": jd.is_clipping,
                "issues": jd.issues,
            }
            for jd in sd.joint_diagnostics
        ],
        "gripper_analysis": {
            "gripper_dim": sd.gripper_analysis.gripper_dim,
            "open_ratio": round(sd.gripper_analysis.open_ratio, 4),
            "close_ratio": round(sd.gripper_analysis.close_ratio, 4),
            "transitions_per_episode": round(sd.gripper_analysis.transitions_per_episode, 2),
            "has_sufficient_grasps": sd.gripper_analysis.has_sufficient_grasps,
            "issues": sd.gripper_analysis.issues,
        },
        "episode_diagnostics": [
            {
                "episode_index": ed.episode_index,
                "quality_score": round(ed.quality_score, 4),
                "length": ed.length,
                "is_truncated": ed.is_truncated,
                "jerk_spike_count": ed.jerk_spike_count,
                "recommendation": ed.recommendation,
                "issues": ed.issues,
            }
            for ed in sd.episode_diagnostics
        ],
        "global_zero_velocity_ratio": round(sd.global_zero_velocity_ratio, 4),
        "dead_dimensions": sd.dead_dimensions,
        "episodes_to_remove": sd.episodes_to_remove,
        "episodes_to_review": sd.episodes_to_review,
        "blockers": sd.blockers,
        "warnings": sd.warnings,
        "passing": sd.passing,
        "prescription": sd.prescription,
    }


def full_report_to_dict(
    info,
    coverage,
    quality,
    assessment,
    recs,
    embedding_analysis=None,
    signal_diagnostics=None,
) -> dict:
    """Convert the full analysis report to a JSON-serializable dict."""
    data = info_to_dict(info)
    data["coverage"] = {
        "position_diversity": coverage.position_diversity,
        "action_diversity": coverage.action_diversity,
        "episode_consistency": coverage.episode_consistency,
        "temporal_coverage": coverage.temporal_coverage,
    }
    data["quality_score"] = quality.overall_score
    data["component_scores"] = quality.component_scores
    data["ready_to_train"] = quality.ready_to_train
    data["ready_to_train_verdict"] = assessment.verdict
    data["ready_to_train_reason"] = assessment.reason
    data["confidence"] = quality.confidence
    low, high = quality.estimated_success_rate
    data["estimated_success_rate"] = {
        "low": round(low, 2),
        "high": round(high, 2),
        "display": f"{low:.0%}-{high:.0%}",
    }
    data["recommendations"] = [
        {
            "priority": r.priority,
            "category": r.category,
            "description": r.description,
            "expected_impact": r.expected_impact,
            "effort": r.effort,
        }
        for r in recs
    ]
    data["blocking_issues"] = assessment.blocking_issues
    data["nice_to_haves"] = assessment.nice_to_haves

    if embedding_analysis is not None:
        data["embedding_analysis"] = {
            "n_clusters": embedding_analysis.n_clusters,
            "cluster_sizes": embedding_analysis.cluster_sizes,
            "visual_diversity": embedding_analysis.diversity_score,
            "outlier_episodes": embedding_analysis.outlier_indices,
            "insights": embedding_analysis.insights,
        }
        data["visual_diversity"] = embedding_analysis.diversity_score

    if signal_diagnostics is not None:
        data["signal_diagnostics"] = signal_diag_to_dict(signal_diagnostics)

    return data
