import io
import json
import zipfile

from riff_pbix import (
    parse_report_layout,
    parse_data_model,
    parse_mashup,
    infer_model_from_references,
    build_lineage,
    generate_plan,
)


def _build_pbix_zip():
    mem = io.BytesIO()
    with zipfile.ZipFile(mem, "w") as z:
        layout = {
            "sections": [
                {
                    "displayName": "Page 1",
                    "name": "section1",
                    "width": 1280,
                    "height": 720,
                    "visualContainers": [
                        {
                            "x": 0,
                            "y": 0,
                            "width": 300,
                            "height": 200,
                            "z": 1,
                            "tabOrder": 0,
                            "config": json.dumps(
                                {
                                    "name": "visual1",
                                    "singleVisual": {
                                        "visualType": "clusteredColumnChart",
                                        "prototypeQuery": {
                                            "From": [{"Entity": "Sales"}],
                                            "Select": [
                                                {"Column": {"Property": "Region"}, "Name": "Region"},
                                                {"Measure": {"Property": "Revenue"}, "Name": "Revenue"},
                                            ],
                                        },
                                        "vcObjects": {
                                            "title": [
                                                {
                                                    "properties": {
                                                        "text": {
                                                            "expr": {
                                                                "Literal": {"Value": "Revenue by Region"}
                                                            }
                                                        }
                                                    }
                                                }
                                            ]
                                        },
                                    },
                                }
                            ),
                        }
                    ],
                }
            ]
        }
        z.writestr("Report/Layout", json.dumps(layout).encode("utf-16-le"))

        model = {
            "model": {
                "tables": [
                    {
                        "name": "Sales",
                        "columns": [{"name": "Region", "dataType": "string"}],
                        "measures": [
                            {"name": "Revenue", "expression": "SUM('Sales'[Amount])"}
                        ],
                        "hierarchies": [],
                    }
                ],
                "relationships": [],
            }
        }
        z.writestr("DataModelSchema", json.dumps(model).encode("utf-16-le"))

        inner_mem = io.BytesIO()
        with zipfile.ZipFile(inner_mem, "w") as inner:
            inner.writestr("Formulas/Section1.m", "shared Sales = let Source = 1 in Source;")
        z.writestr("DataMashup", inner_mem.getvalue())

    mem.seek(0)
    return mem


def test_generate_plan_from_analysis():
    pbix = _build_pbix_zip()
    with zipfile.ZipFile(pbix) as z:
        layout = parse_report_layout(z)
        model = parse_data_model(z)
        mashup = parse_mashup(z)

    if not model or model.get("is_missing"):
        model = infer_model_from_references(layout.get("all_visuals_flat", []))

    lineage = build_lineage(
        layout.get("all_visuals_flat", []),
        model,
        mashup.get("queries", []) if mashup else [],
    )

    analysis_result = {
        "filename": "report.pbix",
        "parsed_layout": layout,
        "parsed_model": model,
        "parsed_mashup": mashup,
        "lineage": lineage,
    }

    plan = generate_plan(analysis_result)

    assert plan.app_name == "report"
    assert len(plan.pages) == 1
    assert plan.pages[0].components[0].props.xAxis == "Region"
    assert plan.pages[0].components[0].props.yAxis == ["Revenue"]
