from app.schemas.visual_lesson import ExtractedChartDataset, ExtractedChartRow
from app.services.composition_compiler import CompositionCompiler


def test_chart_without_extracted_data_falls_back_to_illustrative_values():
    compiler = CompositionCompiler()

    payload = compiler.build("project-1", "give me a bar chart")

    dataset = payload.spec.datasets[0]
    assert "illustrative" in payload.spec.answer_markdown.lower()
    assert payload.spec.assumptions == ["The chart dataset is illustrative."]
    assert len(dataset.rows) > 0


def test_chart_uses_grounded_rows_with_categorical_labels():
    compiler = CompositionCompiler()
    extracted = ExtractedChartDataset(
        has_data=True,
        chart_family="bar",
        x_label="Method",
        y_label="Multi-subject average",
        rows=[
            ExtractedChartRow(label="REINFORCE", value=22.6),
            ExtractedChartRow(label="RLOO (ours)", value=31.2),
        ],
    )

    payload = compiler.build("project-1", "I need a plot for visualizing performance difference", extracted=extracted)

    dataset = payload.spec.datasets[0]
    chart = payload.spec.panels[0].chart
    assert [row["x"] for row in dataset.rows] == ["REINFORCE", "RLOO (ours)"]
    assert [row["y"] for row in dataset.rows] == [22.6, 31.2]
    assert chart.x_label == "Method"
    assert chart.y_label == "Multi-subject average"
    x_encoding = next(e for e in chart.encodings if e.channel == "x")
    assert x_encoding.value_type == "nominal"
    assert payload.spec.assumptions == []
    assert "transcribed from your selected context" in payload.spec.answer_markdown
    assert "grounded in selected context" in payload.spec.provenance.interpretation


def test_chart_uses_grounded_rows_with_numeric_labels_as_quantitative():
    compiler = CompositionCompiler()
    extracted = ExtractedChartDataset(
        has_data=True,
        chart_family="line",
        rows=[ExtractedChartRow(label="2020", value=10.0), ExtractedChartRow(label="2021", value=15.0)],
    )

    payload = compiler.build("project-1", "line chart of results", extracted=extracted)

    chart = payload.spec.panels[0].chart
    assert chart.family == "line"
    x_encoding = next(e for e in chart.encodings if e.channel == "x")
    assert x_encoding.value_type == "quantitative"
    dataset = payload.spec.datasets[0]
    assert dataset.rows[0]["x"] == 2020.0


def test_grounded_categorical_data_does_not_default_to_line_when_prompt_has_no_chart_keyword():
    """Regression test: a prompt like 'I need a table for this plot' or 'plot form for easy
    visualization' contains no bar/line/scatter keyword, so the old keyword-only heuristic always
    fell through to 'line' even for a table of unordered named categories."""
    compiler = CompositionCompiler()
    extracted = ExtractedChartDataset(
        has_data=True,
        chart_family="bar",
        rows=[
            ExtractedChartRow(label="Basic Medicine", value=9.9),
            ExtractedChartRow(label="Law", value=9.2),
            ExtractedChartRow(label="Economics", value=6.7),
        ],
    )

    payload = compiler.build("project-1", "i need a table for this plot", extracted=extracted)

    assert payload.spec.panels[0].chart.family == "bar"


def test_explicit_force_graph_keyword_overrides_extracted_chart_family():
    compiler = CompositionCompiler()
    extracted = ExtractedChartDataset(has_data=True, chart_family="bar", rows=[ExtractedChartRow(label="A", value=1.0)])

    payload = compiler.build("project-1", "force graph of the network", extracted=extracted)

    assert payload.spec.panels[0].chart.family == "force-graph"
    assert payload.spec.datasets[0].columns == ["source", "target", "weight"]


def test_build_routes_to_chart_when_extracted_data_present_without_chart_keyword():
    compiler = CompositionCompiler()
    extracted = ExtractedChartDataset(
        has_data=True, rows=[ExtractedChartRow(label="A", value=1.0), ExtractedChartRow(label="B", value=2.0)]
    )

    payload = compiler.build("project-1", "visualize the performance difference between methods", extracted=extracted)

    assert len(payload.spec.panels) == 1
    assert payload.spec.panels[0].chart.family in {"line", "bar", "scatter", "heatmap"}