"""Tests for crystallise.common.export."""
import pandas as pd
import openpyxl
import io

from crystallise.common.export import to_excel_bytes


class TestToExcelBytes:
    def test_returns_bytes(self):
        results_df = pd.DataFrame({"col": [1, 2]})
        stats = {"total": 2, "errors": 0}
        prompt_info = {"model": "gpt-4.1"}
        original_df = pd.DataFrame({"ID": ["a", "b"]})

        data = to_excel_bytes(results_df, stats, prompt_info, original_df)
        assert isinstance(data, bytes)
        assert len(data) > 0

    def test_creates_expected_sheets(self):
        results_df = pd.DataFrame({"study_type": ["RCT"]})
        stats = {"total": 1}
        prompt_info = {"model": "gpt-4.1"}
        original_df = pd.DataFrame({"ID": ["1"], "Title": ["A"]})
        schema_df = pd.DataFrame({"name": ["study_type"], "type": ["string"]})

        data = to_excel_bytes(results_df, stats, prompt_info, original_df, template_schema=schema_df)

        wb = openpyxl.load_workbook(io.BytesIO(data))
        sheet_names = wb.sheetnames
        assert "Indexing Results" in sheet_names
        assert "Original Data" in sheet_names
        assert "Statistics" in sheet_names
        assert "Configuration" in sheet_names
        assert "Template Schema" in sheet_names

    def test_without_template_schema(self):
        results_df = pd.DataFrame({"col": [1]})
        stats = {"total": 1}
        prompt_info = {"model": "gpt-4.1"}
        original_df = pd.DataFrame({"ID": ["1"]})

        data = to_excel_bytes(results_df, stats, prompt_info, original_df)
        wb = openpyxl.load_workbook(io.BytesIO(data))
        assert "Template Schema" not in wb.sheetnames