"""Tests for data export framework."""

from __future__ import annotations

import csv
import io
import json

import pytest

from orxaq.fabric.exporters import (
    DAGCSVExporter,
    ResultExporter,
    export_dag_summary,
    export_ecl_results,
)
from orxaq.kernel.types import CausalEdge, CausalRole, Entity, VariableType
from orxaq.kernel.world_dag import WorldDAG


def _make_dag() -> WorldDAG:
    """Build a small test DAG with 3 nodes and 2 edges."""
    dag = WorldDAG()
    dag.add_entity(
        Entity(
            entity_id="e1",
            name="Income",
            causal_role=CausalRole.TREATMENT,
            variable_type=VariableType.DYNAMIC,
            domain="credit",
        )
    )
    dag.add_entity(
        Entity(
            entity_id="e2",
            name="Score",
            causal_role=CausalRole.OUTCOME,
            variable_type=VariableType.DYNAMIC,
            domain="credit",
        )
    )
    dag.add_entity(
        Entity(
            entity_id="e3",
            name="Default",
            causal_role=CausalRole.OUTCOME,
            variable_type=VariableType.STATIC,
            domain="credit",
        )
    )
    dag.add_edge(
        CausalEdge(
            source="e1",
            target="e2",
            mechanism="linear",
            confidence=0.85,
            soliton_score=0.9,
        )
    )
    dag.add_edge(
        CausalEdge(
            source="e2",
            target="e3",
            mechanism="threshold",
            confidence=0.75,
            soliton_score=0.8,
        )
    )
    return dag


class TestDAGCSVExporterEdges:
    def test_header_row(self):
        dag = _make_dag()
        result = DAGCSVExporter.export_edges(dag)
        reader = csv.reader(io.StringIO(result))
        header = next(reader)
        assert header == ["source", "target", "mechanism", "confidence", "soliton_score"]

    def test_edge_count(self):
        dag = _make_dag()
        result = DAGCSVExporter.export_edges(dag)
        rows = list(csv.reader(io.StringIO(result)))
        assert len(rows) == 3  # header + 2 edges

    def test_edge_values(self):
        dag = _make_dag()
        result = DAGCSVExporter.export_edges(dag)
        rows = list(csv.reader(io.StringIO(result)))
        # Sorted by (source, target): e1->e2 first, then e2->e3
        assert rows[1][0] == "e1"
        assert rows[1][1] == "e2"
        assert rows[1][2] == "linear"
        assert float(rows[1][3]) == pytest.approx(0.85)

    def test_empty_dag(self):
        dag = WorldDAG()
        result = DAGCSVExporter.export_edges(dag)
        rows = list(csv.reader(io.StringIO(result)))
        assert len(rows) == 1  # header only


class TestDAGCSVExporterNodes:
    def test_header_row(self):
        dag = _make_dag()
        result = DAGCSVExporter.export_nodes(dag)
        reader = csv.reader(io.StringIO(result))
        header = next(reader)
        assert header == ["entity_id", "name", "causal_role", "variable_type", "domain"]

    def test_node_count(self):
        dag = _make_dag()
        result = DAGCSVExporter.export_nodes(dag)
        rows = list(csv.reader(io.StringIO(result)))
        assert len(rows) == 4  # header + 3 nodes

    def test_node_values(self):
        dag = _make_dag()
        result = DAGCSVExporter.export_nodes(dag)
        rows = list(csv.reader(io.StringIO(result)))
        # Sorted by entity_id: e1, e2, e3
        assert rows[1][0] == "e1"
        assert rows[1][1] == "Income"
        assert rows[1][2] == "treatment"

    def test_empty_dag(self):
        dag = WorldDAG()
        result = DAGCSVExporter.export_nodes(dag)
        rows = list(csv.reader(io.StringIO(result)))
        assert len(rows) == 1  # header only


class TestResultExporter:
    def test_to_json_basic(self):
        records = [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]
        result = ResultExporter.to_json(records)
        parsed = json.loads(result)
        assert len(parsed) == 2
        assert parsed[0]["a"] == 1

    def test_to_json_empty(self):
        result = ResultExporter.to_json([])
        assert json.loads(result) == []

    def test_to_csv_basic(self):
        records = [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]
        result = ResultExporter.to_csv(records)
        rows = list(csv.reader(io.StringIO(result)))
        assert rows[0] == ["name", "age"]
        assert rows[1][0] == "Alice"
        assert rows[2][0] == "Bob"

    def test_to_csv_empty(self):
        result = ResultExporter.to_csv([])
        assert result == ""

    def test_to_csv_missing_key(self):
        records = [{"a": 1, "b": 2}, {"a": 3}]
        result = ResultExporter.to_csv(records)
        rows = list(csv.reader(io.StringIO(result)))
        assert rows[2] == ["3", ""]  # missing 'b' key gets empty string

    def test_write_file(self, tmp_path):
        content = "test data\nline 2"
        path = ResultExporter.write(content, tmp_path / "out.txt")
        assert path.exists()
        assert path.read_text() == content

    def test_write_creates_parent_dirs(self, tmp_path):
        path = ResultExporter.write("data", tmp_path / "sub" / "dir" / "file.csv")
        assert path.exists()
        assert path.read_text() == "data"


class TestExportDagSummary:
    def test_returns_list_of_dicts(self):
        dag = _make_dag()
        records = export_dag_summary(dag)
        assert isinstance(records, list)
        assert len(records) == 2  # 2 edges

    def test_record_keys(self):
        dag = _make_dag()
        records = export_dag_summary(dag)
        expected_keys = {
            "source",
            "source_name",
            "target",
            "target_name",
            "mechanism",
            "confidence",
            "soliton_score",
            "edge_type",
        }
        assert set(records[0].keys()) == expected_keys

    def test_source_name_resolved(self):
        dag = _make_dag()
        records = export_dag_summary(dag)
        # First edge e1->e2
        assert records[0]["source"] == "e1"
        assert records[0]["source_name"] == "Income"
        assert records[0]["target"] == "e2"
        assert records[0]["target_name"] == "Score"

    def test_empty_dag(self):
        dag = WorldDAG()
        records = export_dag_summary(dag)
        assert records == []


class TestExportEclResults:
    def test_passthrough_list(self):
        data = [{"segment": "test", "ecl": 100}]
        result = export_ecl_results(data)
        assert result is data

    def test_duck_typed_result(self):
        class FakeResult:
            total_ecl = 1234.56
            ecl_by_segment = {"low_risk": 100.0, "high_risk": 1134.56}

        records = export_ecl_results(FakeResult())
        assert len(records) == 3  # 2 segments + TOTAL
        assert records[-1]["segment"] == "TOTAL"
        assert records[-1]["ecl"] == "1234.5600"

    def test_unknown_type_returns_empty(self):
        result = export_ecl_results("not a list or engine")
        assert result == []
