import json

import numpy as np
import wise_pizza

from wise_pizza_mcp.extract import (
    extract_analysis,
    extract_result,
    extract_segment,
    result_to_json,
)
from wise_pizza_mcp.models import AnalysisResult, SegmentResult, SplitAnalysisResult


def test_extract_segment_basic():
    """Verify SegmentResult conversion from a raw segment dict."""
    raw = {
        "segment": {"region": "US", "product": "A"},
        "index": 0,
        "orig_i": 0,
        "total": np.float64(100.5),
        "seg_size": np.float64(50.0),
        "naive_avg": np.float64(2.01),
        "impact": np.float64(10.0),
        "avg_impact": np.float64(0.5),
        "dummy": np.array([1, 0, 1]),
    }
    result = extract_segment(raw)
    assert isinstance(result, SegmentResult)
    assert result.segment == {"region": "US", "product": "A"}
    assert isinstance(result.total, float)
    assert result.total == 100.5
    assert result.impact == 10.0
    assert result.avg_impact == 0.5


def test_extract_segment_no_impact():
    """Tree solver segments lack impact fields."""
    raw = {
        "segment": {"region": "EU"},
        "index": 1,
        "orig_i": 1,
        "total": 200.0,
        "seg_size": 100.0,
        "naive_avg": 2.0,
        "dummy": np.array([0, 1]),
    }
    result = extract_segment(raw)
    assert result.impact is None
    assert result.avg_impact is None


def test_extract_analysis(sample_df):
    """Run explain_levels on fixture data and verify AnalysisResult fields."""
    sf = wise_pizza.explain_levels(
        df=sample_df,
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
        max_depth=2,
        solver="lasso",
    )
    result = extract_analysis(sf)
    assert isinstance(result, AnalysisResult)
    assert result.task == "levels"
    assert len(result.segments) > 0
    for seg in result.segments:
        assert isinstance(seg.segment, dict)
        assert isinstance(seg.total, float)


def test_extract_slicerpair(sample_df_pair):
    """Run explain_changes with split_fits and verify SplitAnalysisResult."""
    df1, df2 = sample_df_pair
    sp = wise_pizza.explain_changes_in_totals(
        df1=df1,
        df2=df2,
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
        how="split_fits",
        solver="lasso",
    )
    result = extract_result(sp)
    assert isinstance(result, SplitAnalysisResult)
    assert isinstance(result.size_analysis, AnalysisResult)
    assert isinstance(result.average_analysis, AnalysisResult)
    assert len(result.size_analysis.segments) > 0
    assert len(result.average_analysis.segments) > 0


def test_no_numpy_in_json(sample_df):
    """Verify serialized JSON has no numpy types."""
    sf = wise_pizza.explain_levels(
        df=sample_df,
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
        max_depth=2,
        solver="lasso",
    )
    result = extract_analysis(sf)
    json_str = result_to_json(result)

    # Should parse cleanly as JSON
    parsed = json.loads(json_str)
    assert isinstance(parsed, dict)
    assert "segments" in parsed

    # Walk all values and ensure none are numpy types
    def check_no_numpy(obj):
        if isinstance(obj, dict):
            for v in obj.values():
                check_no_numpy(v)
        elif isinstance(obj, list):
            for item in obj:
                check_no_numpy(item)
        else:
            assert not isinstance(obj, (np.integer, np.floating, np.ndarray)), (
                f"Found numpy type: {type(obj)}"
            )

    check_no_numpy(parsed)
