from __future__ import annotations

from datetime import datetime
from typing import Dict, Iterable, List, Optional, Tuple

import sqlalchemy as sa
from sqlalchemy import select, update
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.orm import Session

from db.models import (
    AttributeDef,
    AttributeDiscovery,
    AttributeOption,
    ImageShortUrl,
    IngestRun,
    MagentoProductSourceSnapshot,
    ProductSnapshot,
    ProductSupplementalAttributeValue,
    PlytixProductSourceSnapshot,
    RawFeedRow,
    ShopifyProductSourceSnapshot,
    SkuOverride,
    SkuPriceSnapshot,
)
from db.schemas import (
    AttributeDiscoveryCreate,
    ImageShortUrlCreate,
    ImageShortUrlRead,
    IngestRunCreate,
    MagentoProductSourceSnapshotCreate,
    ProductSupplementalAttributeValueCreate,
    PlytixProductSourceSnapshotCreate,
    ProductSnapshotClose,
    ProductSnapshotCreate,
    RawFeedRowCreate,
    ShopifyProductSourceSnapshotCreate,
    SkuOverrideUpsert,
)


class SqlAlchemyIngestRepository:
    def __init__(self, session: Session) -> None:
        self._session = session

    def create_ingest_run(self, payload: IngestRunCreate) -> int:
        run = IngestRun(**payload.model_dump())
        self._session.add(run)
        self._session.flush()
        return int(run.id)

    def update_ingest_run(self, run_id: int, *, status: str, row_count: int) -> None:
        self._session.execute(
            update(IngestRun)
            .where(IngestRun.id == run_id)
            .values(status=status, row_count=row_count, finished_at=datetime.utcnow())
        )

    def add_raw_rows(self, rows: Iterable[RawFeedRowCreate]) -> None:
        payloads = [RawFeedRow(**row.model_dump()) for row in rows]
        self._session.add_all(payloads)

    def load_current_snapshot_index(self, skus: List[str]) -> Dict[str, Dict[str, object]]:
        if not skus:
            return {}
        stmt = (
            select(ProductSnapshot.id, ProductSnapshot.sku, ProductSnapshot.row_hash)
            .where(ProductSnapshot.sku.in_(skus))
            .where(ProductSnapshot.valid_to.is_(None))
        )
        return {
            sku: {"id": snap_id, "row_hash": row_hash}
            for snap_id, sku, row_hash in self._session.execute(stmt).all()
        }

    def close_snapshots(self, snapshots: Iterable[ProductSnapshotClose]) -> None:
        for snap in snapshots:
            self._session.execute(
                update(ProductSnapshot)
                .where(ProductSnapshot.id == snap.id)
                .values(valid_to=snap.valid_to)
            )

    def add_product_snapshots(self, snapshots: Iterable[ProductSnapshotCreate]) -> None:
        payloads = [ProductSnapshot(**snap.model_dump()) for snap in snapshots]
        self._session.add_all(payloads)

    def close_magento_source_snapshots(
        self,
        snapshot_ids: Iterable[int],
        *,
        valid_to: datetime,
    ) -> None:
        ids = [int(i) for i in snapshot_ids if i]
        if not ids:
            return
        self._session.execute(
            update(MagentoProductSourceSnapshot)
            .where(MagentoProductSourceSnapshot.id.in_(ids))
            .values(valid_to=valid_to)
        )

    def close_plytix_source_snapshots(
        self,
        snapshot_ids: Iterable[int],
        *,
        valid_to: datetime,
    ) -> None:
        ids = [int(i) for i in snapshot_ids if i]
        if not ids:
            return
        self._session.execute(
            update(PlytixProductSourceSnapshot)
            .where(PlytixProductSourceSnapshot.id.in_(ids))
            .values(valid_to=valid_to)
        )

    def close_shopify_source_snapshots(
        self,
        snapshot_ids: Iterable[int],
        *,
        valid_to: datetime,
    ) -> None:
        ids = [int(i) for i in snapshot_ids if i]
        if not ids:
            return
        self._session.execute(
            update(ShopifyProductSourceSnapshot)
            .where(ShopifyProductSourceSnapshot.id.in_(ids))
            .values(valid_to=valid_to)
        )

    def add_magento_source_snapshots(
        self,
        snapshots: Iterable[MagentoProductSourceSnapshotCreate],
    ) -> None:
        payloads = [MagentoProductSourceSnapshot(**snap.model_dump()) for snap in snapshots]
        if payloads:
            self._session.add_all(payloads)

    def add_plytix_source_snapshots(
        self,
        snapshots: Iterable[PlytixProductSourceSnapshotCreate],
    ) -> None:
        payloads = [PlytixProductSourceSnapshot(**snap.model_dump()) for snap in snapshots]
        if payloads:
            self._session.add_all(payloads)

    def add_shopify_source_snapshots(
        self,
        snapshots: Iterable[ShopifyProductSourceSnapshotCreate],
    ) -> None:
        payloads = [ShopifyProductSourceSnapshot(**snap.model_dump()) for snap in snapshots]
        if payloads:
            self._session.add_all(payloads)

    def add_supplemental_attribute_values(
        self,
        rows: Iterable[ProductSupplementalAttributeValueCreate],
    ) -> None:
        payloads = [ProductSupplementalAttributeValue(**row.model_dump()) for row in rows]
        if payloads:
            self._session.add_all(payloads)

    def load_current_magento_source_index(
        self,
        skus: List[str],
        *,
        connection_id: Optional[int] = None,
    ) -> Dict[str, Dict[str, object]]:
        if not skus:
            return {}
        stmt = (
            select(
                MagentoProductSourceSnapshot.id,
                MagentoProductSourceSnapshot.sku,
                MagentoProductSourceSnapshot.row_hash,
            )
            .where(MagentoProductSourceSnapshot.sku.in_(skus))
            .where(MagentoProductSourceSnapshot.valid_to.is_(None))
        )
        if connection_id is None:
            stmt = stmt.where(MagentoProductSourceSnapshot.connection_id.is_(None))
        else:
            stmt = stmt.where(MagentoProductSourceSnapshot.connection_id == connection_id)
        return {
            sku: {"id": snap_id, "row_hash": row_hash}
            for snap_id, sku, row_hash in self._session.execute(stmt).all()
        }

    def load_current_plytix_source_index(self, skus: List[str]) -> Dict[str, Dict[str, object]]:
        if not skus:
            return {}
        stmt = (
            select(
                PlytixProductSourceSnapshot.id,
                PlytixProductSourceSnapshot.sku,
                PlytixProductSourceSnapshot.row_hash,
            )
            .where(PlytixProductSourceSnapshot.sku.in_(skus))
            .where(PlytixProductSourceSnapshot.valid_to.is_(None))
        )
        return {
            sku: {"id": snap_id, "row_hash": row_hash}
            for snap_id, sku, row_hash in self._session.execute(stmt).all()
        }

    def load_current_shopify_source_index(
        self,
        skus: List[str],
        *,
        connection_id: Optional[int] = None,
    ) -> Dict[str, Dict[str, object]]:
        if not skus:
            return {}
        stmt = (
            select(
                ShopifyProductSourceSnapshot.id,
                ShopifyProductSourceSnapshot.sku,
                ShopifyProductSourceSnapshot.row_hash,
            )
            .where(ShopifyProductSourceSnapshot.sku.in_(skus))
            .where(ShopifyProductSourceSnapshot.valid_to.is_(None))
        )
        if connection_id is None:
            stmt = stmt.where(ShopifyProductSourceSnapshot.connection_id.is_(None))
        else:
            stmt = stmt.where(ShopifyProductSourceSnapshot.connection_id == connection_id)
        return {
            sku: {"id": snap_id, "row_hash": row_hash}
            for snap_id, sku, row_hash in self._session.execute(stmt).all()
        }

    def load_all_current_skus(self) -> List[str]:
        """Return all SKUs that have a current (valid_to IS NULL) snapshot."""
        stmt = (
            select(ProductSnapshot.sku)
            .where(ProductSnapshot.valid_to.is_(None))
        )
        return [row for (row,) in self._session.execute(stmt).all() if row]

    def load_current_prices_for_skus(self, skus: List[str]) -> Dict[str, Dict[str, object]]:
        """Return {sku: {price, qty, is_in_stock}} for current price snapshots."""
        if not skus:
            return {}
        stmt = (
            select(SkuPriceSnapshot.sku, SkuPriceSnapshot.price, SkuPriceSnapshot.qty, SkuPriceSnapshot.is_in_stock)
            .where(SkuPriceSnapshot.sku.in_(skus))
            .where(SkuPriceSnapshot.valid_to.is_(None))
        )
        return {
            sku: {"price": price, "qty": qty, "is_in_stock": is_in_stock}
            for sku, price, qty, is_in_stock in self._session.execute(stmt).all()
        }

    def load_current_snapshots(self, skus: List[str]) -> List[dict]:
        if not skus:
            return []
        stmt = (
            select(ProductSnapshot)
            .where(ProductSnapshot.sku.in_(skus))
            .where(ProductSnapshot.valid_to.is_(None))
        )
        rows = self._session.scalars(stmt).all()
        return [
            {
                "id": row.id,
                "sku": row.sku,
                "name": row.name,
                "description": row.description,
                "short_description": row.short_description,
                "product_type": row.product_type,
                "attribute_set_code": row.attribute_set_code,
                "categories_raw": row.categories_raw,
                "base_image": row.base_image,
                "additional_images": row.additional_images,
                "variant_of": row.variant_of,
                "variant_list": row.variant_list,
                "attrs": row.attrs or {},
            }
            for row in rows
        ]

    def load_children_of_parents(self, parent_skus: List[str]) -> List[str]:
        """Return child SKUs whose variant_of is in parent_skus (current snapshots only)."""
        if not parent_skus:
            return []
        stmt = (
            select(ProductSnapshot.sku)
            .where(ProductSnapshot.variant_of.in_(parent_skus))
            .where(ProductSnapshot.valid_to.is_(None))
        )
        return [row for (row,) in self._session.execute(stmt).all() if row]

    def load_attribute_def_index(self) -> Dict[str, Dict[str, object]]:
        stmt = select(AttributeDef.attribute_code, AttributeDef.type, AttributeDef.allows_free_text)
        rows = self._session.execute(stmt).all()
        return {
            code: {
                "type": attr_type or "",
                "allows_free_text": bool(allows_free_text),
            }
            for code, attr_type, allows_free_text in rows
        }

    def load_attribute_option_index(self) -> Dict[str, List[str]]:
        stmt = select(AttributeOption.attribute_code, AttributeOption.option_value)
        rows = self._session.execute(stmt).all()
        index: Dict[str, List[str]] = {}
        for code, value in rows:
            index.setdefault(code, []).append(value)
        return index

    def add_attribute_discoveries(self, rows: Iterable[AttributeDiscoveryCreate]) -> None:
        payloads = [row.model_dump() for row in rows]
        if not payloads:
            return
        stmt = insert(AttributeDiscovery).values(payloads)
        stmt = stmt.on_conflict_do_nothing(index_elements=["attribute_code", "new_value"])
        self._session.execute(stmt)

    def upsert_overrides(self, rows: Iterable[SkuOverrideUpsert]) -> None:
        payloads = [row.model_dump() for row in rows]
        if not payloads:
            return
        stmt = insert(SkuOverride).values(payloads)
        stmt = stmt.on_conflict_do_update(
            index_elements=["sku", "field"],
            set_={"value": stmt.excluded.value, "source": stmt.excluded.source, "updated_at": sa.func.now()},
        )
        self._session.execute(stmt)

    def delete_overrides(self, items: Iterable[Tuple[str, str]]) -> None:
        items_list = list(items)
        if not items_list:
            return
        for sku, field in items_list:
            self._session.query(SkuOverride).filter(
                SkuOverride.sku == sku, SkuOverride.field == field
            ).delete()

    def load_overrides_for_skus(self, skus: List[str]) -> Dict[str, Dict[str, str]]:
        if not skus:
            return {}
        stmt = select(SkuOverride.sku, SkuOverride.field, SkuOverride.value).where(
            SkuOverride.sku.in_(skus)
        )
        rows = self._session.execute(stmt).all()
        out: Dict[str, Dict[str, str]] = {}
        for sku, field, value in rows:
            out.setdefault(sku, {})[field] = value
        return out

    def get_short_url_by_slug(self, slug: str) -> Optional[ImageShortUrlRead]:
        row = self._session.scalars(
            select(ImageShortUrl).where(ImageShortUrl.slug == slug)
        ).first()
        if row is None:
            return None
        return ImageShortUrlRead(
            id=row.id, slug=row.slug, destination_url=row.destination_url, created_at=row.created_at
        )

    def get_short_url_by_destination(self, destination_url: str) -> Optional[ImageShortUrlRead]:
        row = self._session.scalars(
            select(ImageShortUrl).where(ImageShortUrl.destination_url == destination_url)
        ).first()
        if row is None:
            return None
        return ImageShortUrlRead(
            id=row.id, slug=row.slug, destination_url=row.destination_url, created_at=row.created_at
        )

    def create_short_url(self, payload: ImageShortUrlCreate) -> ImageShortUrlRead:
        row = ImageShortUrl(slug=payload.slug, destination_url=payload.destination_url)
        self._session.add(row)
        self._session.flush()
        self._session.refresh(row)
        return ImageShortUrlRead(
            id=row.id, slug=row.slug, destination_url=row.destination_url, created_at=row.created_at
        )