import json
from typing import Union

import numpy as np

from wise_pizza_mcp.models import AnalysisResult, SegmentResult, SplitAnalysisResult


# Keys to skip when building SegmentResult from a raw segment dict
_SKIP_KEYS = {"dummy", "index", "orig_i", "seg_total_vec", "coef"}


def _to_python_float(v: object) -> float:
    """Convert numpy floats/ints to Python float."""
    if isinstance(v, (np.floating, np.integer)):
        return float(v)
    return float(v)


def extract_segment(seg: dict) -> SegmentResult:
    """Convert a raw SliceFinder segment dict to a SegmentResult."""
    segment_def = {str(k): str(v) for k, v in seg["segment"].items()}
    return SegmentResult(
        segment=segment_def,
        total=_to_python_float(seg["total"]),
        seg_size=_to_python_float(seg["seg_size"]),
        naive_avg=_to_python_float(seg["naive_avg"]),
        impact=_to_python_float(seg["impact"]) if "impact" in seg else None,
        avg_impact=_to_python_float(seg["avg_impact"]) if "avg_impact" in seg else None,
    )


def extract_clusters(sf: object) -> dict:
    """Extract relevant cluster names from a SliceFinder, filtering to actual clusters."""
    raw = sf.relevant_cluster_names
    return {str(k): str(v) for k, v in raw.items() if "_cluster_" in k}


def extract_analysis(sf: object) -> AnalysisResult:
    """Convert a SliceFinder to an AnalysisResult."""
    segments = [extract_segment(s) for s in sf.segments]
    clusters = extract_clusters(sf)

    global_average = None
    if hasattr(sf, "global_average"):
        global_average = _to_python_float(sf.global_average)
    elif hasattr(sf, "reg") and sf.reg is not None and hasattr(sf.reg, "intercept_"):
        global_average = _to_python_float(sf.reg.intercept_)

    pre_total = _to_python_float(sf.pre_total) if hasattr(sf, "pre_total") else None
    post_total = _to_python_float(sf.post_total) if hasattr(sf, "post_total") else None

    markdown_summary = None
    try:
        markdown_summary = sf.markdown_summary
    except Exception:
        pass

    return AnalysisResult(
        task=sf.task,
        segments=segments,
        relevant_clusters=clusters,
        global_average=global_average,
        pre_total=pre_total,
        post_total=post_total,
        markdown_summary=markdown_summary,
    )


def extract_result(result: object) -> Union[AnalysisResult, SplitAnalysisResult]:
    """Convert a SliceFinder or SlicerPair to the appropriate result model.

    SlicerPair.summary() is buggy (it tries to access .segments on the pair
    which doesn't exist), so we extract from .s1 and .s2 directly.
    """
    if hasattr(result, "s1") and hasattr(result, "s2"):
        size_analysis = extract_analysis(result.s1)
        average_analysis = extract_analysis(result.s2)
        return SplitAnalysisResult(
            task=result.task or "changes (split)",
            size_analysis=size_analysis,
            average_analysis=average_analysis,
        )
    return extract_analysis(result)


def result_to_json(result: Union[AnalysisResult, SplitAnalysisResult]) -> str:
    """Serialize an analysis result to JSON string, ensuring no numpy types remain."""
    return result.model_dump_json()
