"""Test loading sample data into Pydantic schema models and saving as Parquet.

Proves the full data flow:
    1. CSV sample data → Pydantic models (Product, Order, Customer, InventoryItem)
    2. Mock Shopify GraphQL nodes → Shopify transformer → Pydantic models
    3. Pydantic models → Parquet files (via polars)
    4. Parquet files → polars DataFrame → verify round-trip
"""

import csv
import tempfile
from datetime import datetime
from decimal import Decimal
from pathlib import Path
from uuid import uuid4

import polars as pl
import pytest

from platoon.data.schema import (
    Customer, InventoryItem, InventoryLevel, LineItem,
    Order, Product, Variant,
)
from platoon.models import deterministic_id

SAMPLES = Path(__file__).parent.parent / "data" / "samples"


# ---------------------------------------------------------------------------
# Helpers: CSV → Pydantic
# ---------------------------------------------------------------------------


def _load_csv(path: Path) -> list[dict]:
    with open(path) as f:
        return list(csv.DictReader(f))


def products_from_csv(rows: list[dict]) -> list[Product]:
    """Map products.csv rows into Product + Variant schema instances."""
    products = []
    for r in rows:
        pid = deterministic_id("products", r["product_id"])
        variant = Variant(
            id=deterministic_id("variants", r["sku"]),
            product_id=pid,
            sku=r["sku"],
            title=r["title"],
            price=Decimal(r["price"]),
            weight=float(r["weight_kg"]),
            source_id=r["sku"],
        )
        product = Product(
            id=pid,
            title=r["title"],
            product_type=r["category"],
            status=r["status"],
            variants=[variant],
            source_id=r["product_id"],
        )
        products.append(product)
    return products


def customers_from_csv(rows: list[dict]) -> list[Customer]:
    """Map customers.csv rows into Customer schema instances."""
    return [
        Customer(
            id=deterministic_id("customers", r["customer_id"]),
            email=r["email"],
            first_name=r["first_name"],
            last_name=r["last_name"],
            created_at=datetime.fromisoformat(r["created_at"]),
            source_id=r["customer_id"],
        )
        for r in rows
    ]


def orders_from_csv(rows: list[dict]) -> list[Order]:
    """Map orders.csv rows into Order + LineItem schema instances.

    orders.csv is flat (one row per line item), so we group by order_id.
    """
    grouped: dict[str, list[dict]] = {}
    for r in rows:
        grouped.setdefault(r["order_id"], []).append(r)

    orders = []
    for order_id, items in grouped.items():
        first = items[0]
        line_items = [
            LineItem(
                id=deterministic_id("line_items", f"{order_id}:{li['product_id']}"),
                product_id=deterministic_id("products", li["product_id"]),
                title=li["product_title"],
                sku=li["sku"],
                quantity=int(li["quantity"]),
                price=Decimal(li["unit_price"]),
                total_discount=Decimal(li["discount"]),
                source_id=f"{order_id}:{li['product_id']}",
            )
            for li in items
        ]
        order = Order(
            id=deterministic_id("orders", order_id),
            order_number=first["order_number"],
            customer_id=deterministic_id("customers", first["customer_id"]),
            email=None,
            financial_status=first["financial_status"],
            fulfillment_status=first["fulfillment_status"] or None,
            total_price=sum(Decimal(li["line_total"]) for li in items),
            total_discounts=sum(Decimal(li["discount"]) for li in items),
            currency=first["currency"],
            line_items=line_items,
            created_at=datetime.fromisoformat(first["order_date"]),
            source_id=order_id,
        )
        orders.append(order)
    return orders


def inventory_from_csv(rows: list[dict]) -> list[InventoryItem]:
    """Map inventory.csv rows into InventoryItem instances (latest snapshot per product)."""
    # Take the latest row per product
    latest: dict[str, dict] = {}
    for r in rows:
        pid = r["product_id"]
        if pid not in latest or r["date"] > latest[pid]["date"]:
            latest[pid] = r

    items = []
    for pid, r in latest.items():
        item_id = deterministic_id("inventory_items", pid)
        level = InventoryLevel(
            id=deterministic_id("inventory_levels", f"{pid}:main"),
            inventory_item_id=item_id,
            location_name="Main Warehouse",
            available=int(r["closing_stock"]),
            on_hand=int(r["closing_stock"]),
        )
        items.append(InventoryItem(
            id=item_id,
            sku=r["sku"],
            cost=Decimal(r["cost_per_unit"]),
            tracked=True,
            levels=[level],
            source_id=pid,
        ))
    return items


# ---------------------------------------------------------------------------
# Helpers: Pydantic → Parquet (via polars)
# ---------------------------------------------------------------------------


def models_to_parquet(models: list, path: Path, flatten_fields: dict | None = None):
    """Serialize Pydantic models to a Parquet file.

    Nested lists (variants, line_items, levels) are excluded from the main table
    and written as separate relationship files.

    Args:
        models: List of Pydantic model instances
        path: Output .parquet path
        flatten_fields: Optional {field_name: child_parquet_path} for nested lists
    """
    # Dump to flat dicts, converting UUIDs and Decimals to strings
    rows = []
    child_rows: dict[str, list] = {k: [] for k in (flatten_fields or {})}

    for m in models:
        d = m.model_dump(mode="json")
        # Extract nested lists to separate tables
        for field_name in (flatten_fields or {}):
            children = d.pop(field_name, [])
            child_rows[field_name].extend(children)
        rows.append(d)

    # Write main table
    df = pl.DataFrame(rows)
    df.write_parquet(path)

    # Write child tables
    for field_name, child_path in (flatten_fields or {}).items():
        if child_rows[field_name]:
            child_df = pl.DataFrame(child_rows[field_name])
            child_df.write_parquet(child_path)

    return df


# ---------------------------------------------------------------------------
# Tests: CSV → Pydantic → Parquet round-trip
# ---------------------------------------------------------------------------


@pytest.fixture
def tmp_dir():
    with tempfile.TemporaryDirectory() as d:
        yield Path(d)


class TestCSVToParquetRoundtrip:
    """Load sample CSVs into Pydantic models and save as Parquet."""

    @pytest.fixture(autouse=True)
    def _check_samples(self):
        if not (SAMPLES / "products.csv").exists():
            pytest.skip("Sample data not generated — run: uv run python data/samples/generate.py")

    def test_products(self, tmp_dir):
        """products.csv → Product + Variant → products.parquet + variants.parquet"""
        rows = _load_csv(SAMPLES / "products.csv")
        products = products_from_csv(rows)

        assert len(products) == 80
        assert all(isinstance(p, Product) for p in products)
        assert all(len(p.variants) == 1 for p in products)

        # Write parquet
        variants_path = tmp_dir / "variants.parquet"
        df = models_to_parquet(
            products,
            tmp_dir / "products.parquet",
            flatten_fields={"variants": variants_path},
        )

        assert (tmp_dir / "products.parquet").exists()
        assert variants_path.exists()

        # Read back and verify
        products_df = pl.read_parquet(tmp_dir / "products.parquet")
        variants_df = pl.read_parquet(variants_path)

        assert len(products_df) == 80
        assert len(variants_df) == 80
        assert "title" in products_df.columns
        assert "product_type" in products_df.columns
        assert "sku" in variants_df.columns
        assert "price" in variants_df.columns

        print(f"\n  Products: {len(products_df)} rows, {products_df.estimated_size()} bytes")
        print(f"  Variants: {len(variants_df)} rows, {variants_df.estimated_size()} bytes")

    def test_customers(self, tmp_dir):
        """customers.csv → Customer → customers.parquet"""
        rows = _load_csv(SAMPLES / "customers.csv")
        customers = customers_from_csv(rows)

        assert len(customers) == 5000
        assert all(isinstance(c, Customer) for c in customers)

        df = models_to_parquet(customers, tmp_dir / "customers.parquet")

        customers_df = pl.read_parquet(tmp_dir / "customers.parquet")
        assert len(customers_df) == 5000
        assert "email" in customers_df.columns
        assert "first_name" in customers_df.columns

        print(f"\n  Customers: {len(customers_df)} rows, {customers_df.estimated_size()} bytes")

    def test_orders(self, tmp_dir):
        """orders.csv → Order + LineItem → orders.parquet + line_items.parquet"""
        rows = _load_csv(SAMPLES / "orders.csv")
        orders = orders_from_csv(rows)

        assert len(orders) > 1000
        assert all(isinstance(o, Order) for o in orders)
        assert all(len(o.line_items) >= 1 for o in orders)

        line_items_path = tmp_dir / "line_items.parquet"
        df = models_to_parquet(
            orders,
            tmp_dir / "orders.parquet",
            flatten_fields={"line_items": line_items_path},
        )

        orders_df = pl.read_parquet(tmp_dir / "orders.parquet")
        line_items_df = pl.read_parquet(line_items_path)

        assert len(orders_df) > 1000
        assert len(line_items_df) > len(orders_df)  # multiple items per order
        assert "total_price" in orders_df.columns
        assert "financial_status" in orders_df.columns
        assert "sku" in line_items_df.columns
        assert "quantity" in line_items_df.columns

        print(f"\n  Orders: {len(orders_df)} rows, {orders_df.estimated_size()} bytes")
        print(f"  Line items: {len(line_items_df)} rows, {line_items_df.estimated_size()} bytes")

    def test_inventory(self, tmp_dir):
        """inventory.csv → InventoryItem + InventoryLevel → parquet"""
        rows = _load_csv(SAMPLES / "inventory.csv")
        items = inventory_from_csv(rows)

        assert len(items) == 80  # one per product
        assert all(isinstance(i, InventoryItem) for i in items)
        assert all(len(i.levels) == 1 for i in items)

        levels_path = tmp_dir / "inventory_levels.parquet"
        df = models_to_parquet(
            items,
            tmp_dir / "inventory_items.parquet",
            flatten_fields={"levels": levels_path},
        )

        items_df = pl.read_parquet(tmp_dir / "inventory_items.parquet")
        levels_df = pl.read_parquet(levels_path)

        assert len(items_df) == 80
        assert len(levels_df) == 80
        assert "sku" in items_df.columns
        assert "cost" in items_df.columns
        assert "available" in levels_df.columns
        assert "on_hand" in levels_df.columns

        print(f"\n  Inventory items: {len(items_df)} rows, {items_df.estimated_size()} bytes")
        print(f"  Inventory levels: {len(levels_df)} rows, {levels_df.estimated_size()} bytes")


# ---------------------------------------------------------------------------
# Tests: Shopify transformer → Pydantic → Parquet
# ---------------------------------------------------------------------------


class TestShopifyToParquet:
    """Shopify GraphQL nodes → transformer → Pydantic → Parquet."""

    def test_shopify_products_to_parquet(self, tmp_dir):
        """Mock Shopify product nodes through the full pipeline."""
        from platoon.data.providers.shopify.transform import ShopifyTransformer

        nodes = [
            {
                "id": f"gid://shopify/Product/{i}",
                "title": f"Product {i}",
                "descriptionHtml": f"<p>Description for product {i}</p>",
                "vendor": "TestVendor",
                "productType": "Widgets",
                "tags": ["test", f"batch-{i % 3}"],
                "status": "ACTIVE",
                "createdAt": "2026-01-15T10:00:00Z",
                "updatedAt": "2026-02-20T14:30:00Z",
                "variants": {"edges": [
                    {"node": {
                        "id": f"gid://shopify/ProductVariant/{i * 10 + j}",
                        "sku": f"WIDGET-{i:03d}-{chr(65 + j)}",
                        "title": f"Variant {chr(65 + j)}",
                        "price": f"{19.99 + i + j * 5:.2f}",
                        "compareAtPrice": f"{29.99 + i:.2f}" if j == 0 else None,
                        "inventoryQuantity": 100 - i * 3 + j * 10,
                        "weight": 0.5 + i * 0.1,
                        "barcode": f"1234567890{i:02d}{j}",
                    }}
                    for j in range(3)  # 3 variants each
                ]},
            }
            for i in range(25)  # 25 products
        ]

        transformer = ShopifyTransformer()
        products = transformer.products(nodes)

        assert len(products) == 25
        assert all(len(p.variants) == 3 for p in products)

        # To parquet
        variants_path = tmp_dir / "shopify_variants.parquet"
        models_to_parquet(
            products,
            tmp_dir / "shopify_products.parquet",
            flatten_fields={"variants": variants_path},
        )

        products_df = pl.read_parquet(tmp_dir / "shopify_products.parquet")
        variants_df = pl.read_parquet(variants_path)

        assert len(products_df) == 25
        assert len(variants_df) == 75  # 25 * 3
        assert products_df["source_id"][0].startswith("gid://shopify/Product/")
        assert "WIDGET-" in variants_df["sku"][0]

        # Deterministic IDs: same input → same UUID
        products2 = transformer.products(nodes)
        assert products[0].id == products2[0].id
        assert products[0].variants[0].id == products2[0].variants[0].id

        print(f"\n  Shopify products: {len(products_df)} rows")
        print(f"  Shopify variants: {len(variants_df)} rows")
        print(f"  Sample source_id: {products_df['source_id'][0]}")

    def test_shopify_orders_to_parquet(self, tmp_dir):
        """Mock Shopify order nodes through the full pipeline."""
        from platoon.data.providers.shopify.transform import ShopifyTransformer

        nodes = [
            {
                "id": f"gid://shopify/Order/{1000 + i}",
                "name": f"#{1000 + i}",
                "email": f"customer{i}@example.com",
                "displayFinancialStatus": "PAID",
                "displayFulfillmentStatus": "FULFILLED",
                "currencyCode": "USD",
                "createdAt": f"2026-02-{(i % 28) + 1:02d}T12:00:00Z",
                "updatedAt": f"2026-02-{(i % 28) + 1:02d}T12:00:00Z",
                "totalPriceSet": {"shopMoney": {"amount": f"{50 + i * 10:.2f}"}},
                "subtotalPriceSet": {"shopMoney": {"amount": f"{45 + i * 10:.2f}"}},
                "totalTaxSet": {"shopMoney": {"amount": f"{5:.2f}"}},
                "totalDiscountsSet": {"shopMoney": {"amount": f"{i * 2:.2f}"}},
                "customer": {"id": f"gid://shopify/Customer/{i % 10}"},
                "lineItems": {"edges": [
                    {"node": {
                        "id": f"gid://shopify/LineItem/{i * 100 + j}",
                        "title": f"Widget {j}",
                        "sku": f"WIDGET-{j:03d}-A",
                        "quantity": j + 1,
                        "originalUnitPriceSet": {"shopMoney": {"amount": f"{19.99 + j * 5:.2f}"}},
                        "totalDiscountSet": {"shopMoney": {"amount": "0.00"}},
                        "product": {"id": f"gid://shopify/Product/{j}"},
                        "variant": {"id": f"gid://shopify/ProductVariant/{j * 10}"},
                    }}
                    for j in range(2)
                ]},
            }
            for i in range(50)
        ]

        transformer = ShopifyTransformer()
        orders = transformer.orders(nodes)

        assert len(orders) == 50
        assert all(len(o.line_items) == 2 for o in orders)

        line_items_path = tmp_dir / "shopify_line_items.parquet"
        models_to_parquet(
            orders,
            tmp_dir / "shopify_orders.parquet",
            flatten_fields={"line_items": line_items_path},
        )

        orders_df = pl.read_parquet(tmp_dir / "shopify_orders.parquet")
        li_df = pl.read_parquet(line_items_path)

        assert len(orders_df) == 50
        assert len(li_df) == 100
        assert "total_price" in orders_df.columns
        assert "financial_status" in orders_df.columns

        print(f"\n  Shopify orders: {len(orders_df)} rows")
        print(f"  Shopify line items: {len(li_df)} rows")


# ---------------------------------------------------------------------------
# Test: Full export to a data directory
# ---------------------------------------------------------------------------


class TestFullExport:
    """Export all entity types to a parquet data directory."""

    @pytest.fixture(autouse=True)
    def _check_samples(self):
        if not (SAMPLES / "products.csv").exists():
            pytest.skip("Sample data not generated — run: uv run python data/samples/generate.py")

    def test_full_export(self, tmp_dir):
        """Load all CSVs → Pydantic → Parquet, print summary."""
        out = tmp_dir / "warehouse"
        out.mkdir()

        # Products
        products = products_from_csv(_load_csv(SAMPLES / "products.csv"))
        models_to_parquet(products, out / "products.parquet",
                          flatten_fields={"variants": out / "variants.parquet"})

        # Customers
        customers = customers_from_csv(_load_csv(SAMPLES / "customers.csv"))
        models_to_parquet(customers, out / "customers.parquet")

        # Orders
        orders = orders_from_csv(_load_csv(SAMPLES / "orders.csv"))
        models_to_parquet(orders, out / "orders.parquet",
                          flatten_fields={"line_items": out / "line_items.parquet"})

        # Inventory
        items = inventory_from_csv(_load_csv(SAMPLES / "inventory.csv"))
        models_to_parquet(items, out / "inventory_items.parquet",
                          flatten_fields={"levels": out / "inventory_levels.parquet"})

        # Verify all files exist
        expected = [
            "products.parquet", "variants.parquet",
            "customers.parquet",
            "orders.parquet", "line_items.parquet",
            "inventory_items.parquet", "inventory_levels.parquet",
        ]
        for fname in expected:
            assert (out / fname).exists(), f"Missing: {fname}"

        # Print summary
        print(f"\n  {'File':<28} {'Rows':>8} {'Size':>10}")
        print("  " + "-" * 48)
        total_bytes = 0
        for fname in expected:
            df = pl.read_parquet(out / fname)
            size = (out / fname).stat().st_size
            total_bytes += size
            print(f"  {fname:<28} {len(df):>8} {size:>9,} B")
        print(f"  {'TOTAL':<28} {'':>8} {total_bytes:>9,} B")

        # Verify we can query across tables
        orders_df = pl.read_parquet(out / "orders.parquet")
        products_df = pl.read_parquet(out / "products.parquet")

        assert len(orders_df) > 0
        assert len(products_df) == 80

        # Revenue by financial status
        revenue = orders_df.group_by("financial_status").agg(
            pl.col("total_price").cast(pl.Float64).sum().alias("revenue"),
            pl.len().alias("count"),
        )
        print(f"\n  Revenue by status:")
        for row in revenue.to_dicts():
            print(f"    {row['financial_status']:<20} ${row['revenue']:>12,.2f}  ({row['count']} orders)")
