"""Data pipeline orchestration — runs Source -> Transform -> Sink per tenant.

The pipeline manages sync state (watermarks) and coordinates providers
to incrementally load data from external sources into local storage.

Supports two storage backends:
    1. In-memory (default) — for testing and single-run scripts
    2. JSON file — persists state and records to disk

For production use with percolate/Postgres, extend with a database sink.

Usage:
    pipeline = Pipeline()
    await pipeline.run_entity("tenant-1", "shopify", "products", provider=shopify_provider)

    # Or register providers and run everything
    pipeline.register("shopify", shopify_provider)
    await pipeline.run_all("tenant-1")
"""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional
from uuid import UUID, uuid4

from pydantic import BaseModel, Field

from platoon.data.provider import DataProvider, DataProviderResult

logger = logging.getLogger(__name__)


class SyncState(BaseModel):
    """Tracks watermark for a specific tenant + integration + entity_type.

    Allows incremental syncs by recording the last successful sync
    timestamp and stats.

    Example:
        tenant_id=abc, integration=shopify, entity_type=orders,
        last_watermark=2026-02-25T00:00:00, last_count=142, status=success
    """

    id: UUID = Field(default_factory=uuid4)
    tenant_id: str
    integration_name: str
    entity_type: str
    last_watermark: Optional[datetime] = None
    last_run_at: Optional[datetime] = None
    last_count: int = 0
    total_synced: int = 0
    status: str = "pending"  # pending | running | success | error
    error_message: Optional[str] = None


class SyncSummary(BaseModel):
    """Summary of a pipeline run."""

    tenant_id: str
    results: Dict[str, Dict[str, int]] = Field(default_factory=dict)  # {integration: {entity: count}}
    errors: List[str] = Field(default_factory=list)
    started_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    completed_at: Optional[datetime] = None


class Pipeline:
    """Orchestrates data sync from external providers to storage.

    Responsibilities:
        1. Manage sync state (watermarks per tenant/integration/entity)
        2. Route to the correct DataProvider for each integration
        3. Run fetch -> transform -> sink for each entity type
        4. Record success/failure and update watermarks

    Storage modes:
        - In-memory: all state and records live in dicts (default)
        - File-backed: state persists to a JSON file on disk
    """

    def __init__(self, state_file: Optional[str] = None):
        """Initialize pipeline.

        Args:
            state_file: Optional path to persist sync state as JSON.
                        If None, state is in-memory only.
        """
        self._providers: Dict[str, DataProvider] = {}
        self._states: Dict[str, SyncState] = {}  # key = "tenant:integration:entity"
        self._records: Dict[str, List[Any]] = {}  # key = "tenant:entity" -> list of records
        self._state_file = Path(state_file) if state_file else None

        if self._state_file and self._state_file.exists():
            self._load_state()

    def register(self, name: str, provider: DataProvider) -> None:
        """Register a data provider by integration name.

        Args:
            name: Integration name (e.g. "shopify", "stripe")
            provider: DataProvider instance
        """
        self._providers[name] = provider

    def get_provider(self, name: str) -> DataProvider:
        """Get a registered provider by name."""
        if name not in self._providers:
            raise ValueError(f"No provider registered as '{name}'. Available: {list(self._providers.keys())}")
        return self._providers[name]

    async def run_all(self, tenant_id: str) -> SyncSummary:
        """Run sync for a tenant across all registered providers.

        Args:
            tenant_id: Tenant identifier

        Returns:
            SyncSummary with counts and any errors
        """
        summary = SyncSummary(tenant_id=tenant_id)

        for integration_name, provider in self._providers.items():
            summary.results[integration_name] = {}
            for entity_type in provider.supported_entities():
                try:
                    count = await self.run_entity(
                        tenant_id, integration_name, entity_type, provider=provider
                    )
                    summary.results[integration_name][entity_type] = count
                except Exception as e:
                    error = f"{integration_name}/{entity_type}: {e}"
                    summary.errors.append(error)
                    logger.error(error)

        summary.completed_at = datetime.now(timezone.utc)
        return summary

    async def run_entity(
        self,
        tenant_id: str,
        integration_name: str,
        entity_type: str,
        provider: Optional[DataProvider] = None,
        force_full: bool = False,
    ) -> int:
        """Run sync for one specific entity type.

        This is the core unit of work:
            1. Load sync state (watermark)
            2. Fetch from provider since watermark
            3. Sink records to storage
            4. Update watermark

        Args:
            tenant_id: Tenant identifier
            integration_name: e.g. "shopify"
            entity_type: e.g. "orders", "products"
            provider: DataProvider instance (or looks up registered provider)
            force_full: If True, ignores watermark and does a full sync

        Returns:
            Number of records synced
        """
        if provider is None:
            provider = self.get_provider(integration_name)

        state = self.get_sync_state(tenant_id, integration_name, entity_type)
        state.status = "running"
        state.last_run_at = datetime.now(timezone.utc)

        since = None if force_full else state.last_watermark

        try:
            result = await provider.fetch(entity_type, since=since)
            self._sink(tenant_id, entity_type, result.records)

            state.last_watermark = result.watermark
            state.last_count = result.count
            state.total_synced += result.count
            state.status = "success"
            state.error_message = None

            logger.info(
                f"Synced {result.count} {entity_type} for {tenant_id}/{integration_name}"
            )

        except Exception as e:
            state.status = "error"
            state.error_message = str(e)
            logger.error(f"Sync failed {tenant_id}/{integration_name}/{entity_type}: {e}")
            raise

        finally:
            self._save_sync_state(state)

        return state.last_count

    def get_sync_state(
        self, tenant_id: str, integration_name: str, entity_type: str
    ) -> SyncState:
        """Load or create sync state for a tenant/integration/entity combo."""
        key = f"{tenant_id}:{integration_name}:{entity_type}"
        if key not in self._states:
            self._states[key] = SyncState(
                tenant_id=tenant_id,
                integration_name=integration_name,
                entity_type=entity_type,
            )
        return self._states[key]

    def get_records(self, tenant_id: str, entity_type: str) -> List[Any]:
        """Get all synced records for a tenant/entity. Useful for testing."""
        key = f"{tenant_id}:{entity_type}"
        return self._records.get(key, [])

    def list_states(self, tenant_id: Optional[str] = None) -> List[SyncState]:
        """List all sync states, optionally filtered by tenant."""
        states = list(self._states.values())
        if tenant_id:
            states = [s for s in states if s.tenant_id == tenant_id]
        return states

    # ----- Storage -----

    def _sink(self, tenant_id: str, entity_type: str, records: list) -> None:
        """Store records in the in-memory record store.

        For production, override this to upsert into Postgres/percolate.
        Records are Pydantic model instances — serialize via .model_dump().
        """
        key = f"{tenant_id}:{entity_type}"
        if key not in self._records:
            self._records[key] = []

        # Upsert by id: replace existing records with same id
        existing_ids = {r.id for r in self._records[key] if hasattr(r, "id")}
        new_records = []
        for record in records:
            record_id = getattr(record, "id", None)
            if record_id and record_id in existing_ids:
                # Replace existing
                self._records[key] = [
                    r for r in self._records[key]
                    if getattr(r, "id", None) != record_id
                ]
            new_records.append(record)

        self._records[key].extend(new_records)

    def _save_sync_state(self, state: SyncState) -> None:
        """Persist sync state."""
        key = f"{state.tenant_id}:{state.integration_name}:{state.entity_type}"
        self._states[key] = state

        if self._state_file:
            self._persist_state()

    def _persist_state(self) -> None:
        """Write all sync states to the state file."""
        if not self._state_file:
            return
        data = {
            key: state.model_dump(mode="json")
            for key, state in self._states.items()
        }
        self._state_file.parent.mkdir(parents=True, exist_ok=True)
        self._state_file.write_text(json.dumps(data, indent=2, default=str))

    def _load_state(self) -> None:
        """Load sync states from the state file."""
        if not self._state_file or not self._state_file.exists():
            return
        try:
            data = json.loads(self._state_file.read_text())
            for key, state_dict in data.items():
                self._states[key] = SyncState.model_validate(state_dict)
            logger.info(f"Loaded {len(self._states)} sync states from {self._state_file}")
        except Exception as e:
            logger.warning(f"Failed to load sync state: {e}")
