"""
test3.py — Simple test for bribrain load_to_excel()
"""

from unittest.mock import MagicMock, patch
import pandas as pd
import os

from bribrain.load import load_to_excel

# ── Dummy DataFrames (mock PySpark) ──────────────────────────────────────

def make_dummy_df(data: dict) -> MagicMock:
    """Buat mock PySpark DataFrame dari dict."""
    pdf = pd.DataFrame(data)
    mock_df = MagicMock()
    mock_df.coalesce.return_value.toPandas.return_value = pdf
    return mock_df

# Dataframe 1: Data Penjualan
df_sales = make_dummy_df({
    "produk": ["Laptop", "Mouse", "Keyboard"],
    "qty": [10, 50, 30],
    "harga": [15_000_000, 250_000, 750_000],
})

# Dataframe 2: Data Pelanggan
df_customers = make_dummy_df({
    "customer_id": [1, 2, 3],
    "nama": ["Alice", "Bob", "Charlie"],
    "kota": ["Jakarta", "Bandung", "Surabaya"],
})

# Dataframe 3: Data Stok
df_stock = make_dummy_df({
    "kode_barang": ["A001", "A002", "A003"],
    "nama_barang": ["Laptop", "Mouse", "Keyboard"],
    "stok": [100, 500, 300],
})

# ── Test ─────────────────────────────────────────────────────────────────

# Mock subprocess (HDFS) dan download agar bisa jalan tanpa cluster
@patch("bribrain.load.download_hdfs_file")
@patch("subprocess.run", return_value=MagicMock(returncode=0))
@patch.dict(os.environ, {"HADOOP_USER_NAME": "testuser"})
def run_tests(mock_subproc, mock_download):

    # --- Test 1: Single DataFrame, default sheet name ---
    print("=" * 60)
    print("TEST 1: Single DataFrame → default sheet name")
    print("=" * 60)
    with patch("os.remove"):
        load_to_excel(df_sales, "test_single", download=False)

    # Cek file xlsx
    assert os.path.exists("test_single.xlsx"), "File tidak ditemukan!"
    df = pd.read_excel("test_single.xlsx")
    print(df)
    print(f"✅ Sheet names: {pd.ExcelFile('test_single.xlsx').sheet_names}")
    os.remove("test_single.xlsx")
    print()

    # --- Test 2: Multiple DataFrames, custom sheet names ---
    print("=" * 60)
    print("TEST 2: Multiple DataFrames → custom sheet names")
    print("=" * 60)
    with patch("os.remove"):
        load_to_excel(
            [df_sales, df_customers, df_stock],
            "test_multi",
            sheet_name=["Penjualan", "Pelanggan", "Stok"],
            download=False,
        )

    assert os.path.exists("test_multi.xlsx"), "File tidak ditemukan!"
    xls = pd.ExcelFile("test_multi.xlsx")
    print(f"✅ Sheet names: {xls.sheet_names}")
    for sheet in xls.sheet_names:
        print(f"\n--- {sheet} ---")
        print(pd.read_excel(xls, sheet_name=sheet))
    xls.close()
    os.remove("test_multi.xlsx")
    print()

    # --- Test 3: Multiple DataFrames, default sheet names ---
    print("=" * 60)
    print("TEST 3: Multiple DataFrames → default sheet names (Sheet1, Sheet2)")
    print("=" * 60)
    with patch("os.remove"):
        load_to_excel(
            [df_sales, df_customers],
            "test_default_sheets",
            download=False,
        )

    assert os.path.exists("test_default_sheets.xlsx"), "File tidak ditemukan!"
    xls = pd.ExcelFile("test_default_sheets.xlsx")
    print(f"✅ Sheet names: {xls.sheet_names}")
    for sheet in xls.sheet_names:
        print(f"\n--- {sheet} ---")
        print(pd.read_excel(xls, sheet_name=sheet))
    xls.close()
    os.remove("test_default_sheets.xlsx")

    print()
    print("🎉 Semua test PASSED!")

if __name__ == "__main__":
    run_tests()
