#!/usr/bin/env python3
"""
Two-Tier Event Store - Pragmatic LEM Storage

Hot Tier (SQLite): Last 7 days, fast operational queries
Cold Tier (Parquet): Everything, cheap analytics with DuckDB

Design Philosophy:
- Don't use Kafka until you need it (>10M events/day)
- SQLite + WAL mode = bulletproof for millions of rows
- Parquet + DuckDB = billions of rows, pennies per TB
- Partition by date/tenant/type for efficient queries

Usage:
    store = EventStore(data_dir=Path("data"))

    # Write event
    store.write(event)

    # Query recent (hot tier - fast)
    events = store.query_recent(tenant="acme", hours=24)

    # Query historical (cold tier - powerful)
    df = store.query_analytical("SELECT * FROM events WHERE year=2025 AND month=11")
"""

import sqlite3
import json
import logging
from pathlib import Path
from datetime import datetime, timezone, timedelta
from typing import List, Dict, Any, Optional
from dataclasses import asdict

try:
    import duckdb
    HAS_DUCKDB = True
except ImportError:
    HAS_DUCKDB = False
    logging.warning("duckdb not available, analytical queries disabled")

try:
    import pyarrow as pa
    import pyarrow.parquet as pq
    HAS_PARQUET = True
except ImportError:
    HAS_PARQUET = False
    logging.warning("pyarrow not available, cold storage disabled")

from alerts.tiering import AlertTier
from receipts.base_event import BaseEvent, ReceiptEvent, AlertEvent
from receipts.storage_router import route_by_tier, sign_if_needed, generate_court_packet

logger = logging.getLogger(__name__)


class EventStore:
    """
    Pragmatic two-tier event store.

    Hot Tier: SQLite (last 7 days) for operational queries
    Cold Tier: Parquet (everything) for analytics

    Features:
    - Automatic partitioning by date/tenant/type
    - Batch writes to Parquet (every hour)
    - Time-travel queries via DuckDB
    - Multi-tenant isolation
    """

    def __init__(
        self,
        data_dir: Path,
        hot_retention_days: int = 7,
        batch_size: int = 1000,
    ):
        """
        Initialize event store.

        Args:
            data_dir: Root directory for storage
            hot_retention_days: How long to keep in SQLite
            batch_size: Events to buffer before Parquet write
        """
        self.data_dir = Path(data_dir)
        self.hot_retention_days = hot_retention_days
        self.batch_size = batch_size

        # Create directories
        self.hot_dir = self.data_dir / "hot"
        self.cold_dir = self.data_dir / "cold" / "events"
        self.hot_dir.mkdir(parents=True, exist_ok=True)
        self.cold_dir.mkdir(parents=True, exist_ok=True)

        # Initialize SQLite (hot tier)
        self._init_sqlite()

        # Initialize DuckDB (for cold tier queries)
        if HAS_DUCKDB:
            self.duckdb_conn = duckdb.connect(str(self.data_dir / "analytics.duckdb"))
        else:
            self.duckdb_conn = None

        # Write buffer for Parquet batching
        self._write_buffer: List[Dict[str, Any]] = []

        logger.info(f"EventStore initialized: {self.data_dir}")

    def _init_sqlite(self):
        """Initialize SQLite with WAL mode for performance."""
        db_path = self.hot_dir / "events.db"

        self.sqlite_conn = sqlite3.connect(
            str(db_path),
            isolation_level=None,  # Autocommit mode
            check_same_thread=False,  # Allow multi-threaded access
        )

        # Enable WAL mode (Write-Ahead Logging) for better concurrency
        self.sqlite_conn.execute("PRAGMA journal_mode=WAL")
        self.sqlite_conn.execute("PRAGMA synchronous=NORMAL")
        self.sqlite_conn.execute("PRAGMA cache_size=-64000")  # 64MB cache

        # Create events table
        self.sqlite_conn.execute("""
            CREATE TABLE IF NOT EXISTS events (
                event_id TEXT PRIMARY KEY,
                event_type TEXT NOT NULL,
                tenant_id TEXT NOT NULL,
                valid_time REAL NOT NULL,
                transaction_time REAL NOT NULL,
                observed_at REAL,
                dip_hash TEXT,
                signature TEXT,
                parent_ids TEXT,
                namespaces TEXT,
                data JSON NOT NULL,
                privacy_tags TEXT,
                source TEXT,
                tier TEXT,
                chi_score REAL,
                dignity_score REAL,
                traceparent TEXT,
                tracestate TEXT,

                -- Receipt-specific fields
                verdict TEXT,
                cost JSON,
                latency_ms REAL,
                refusal_reason TEXT,

                -- Alert-specific fields
                severity TEXT,
                alert_type TEXT,
                conditions JSON,
                recommended_actions TEXT
            )
        """)

        # Indexes for fast queries
        self.sqlite_conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_tenant_time
            ON events(tenant_id, valid_time DESC)
        """)

        self.sqlite_conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_type_time
            ON events(event_type, valid_time DESC)
        """)

        self.sqlite_conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_parent
            ON events(parent_ids)
        """)

        logger.info("SQLite hot tier initialized")

    def write(self, event: BaseEvent):
        """
        Tier-aware write path routing events to the appropriate backend.
        """
        backend, needs_signing = route_by_tier(event)

        if needs_signing:
            event = sign_if_needed(event)

        if backend == "memory":
            logger.debug("Dropping transient event %s (tier=%s)", event.event_id, event.tier.name)
            return

        if backend == "sqlite_hot":
            self._write_to_sqlite(event)
        elif backend == "parquet_cold":
            # Mirror to hot tier for short-lived queries, archive in cold tier for retention
            self._write_to_sqlite(event)
            self._buffer_cold(event)
        else:
            # Safety: store in both tiers if router returns unknown backend
            self._write_to_sqlite(event)
            self._buffer_cold(event)

        if event.tier == AlertTier.LEGAL:
            try:
                generate_court_packet(event)
            except Exception:  # pragma: no cover - best effort
                logger.exception("Failed to generate CourtPacket for %s", event.event_id)

    def _write_to_sqlite(self, event: BaseEvent):
        """Write single event to SQLite."""
        event_dict = event.to_dict()

        # Convert lists/dicts to JSON strings for SQLite
        parent_ids_json = json.dumps(event.parent_ids)
        namespaces_json = json.dumps(event.namespaces)
        privacy_tags_json = json.dumps(event.privacy_tags)
        data_json = json.dumps(event.data)

        # Convert datetime to timestamp
        valid_ts = event.valid_time.timestamp()
        trans_ts = event.transaction_time.timestamp()
        obs_ts = event.observed_at.timestamp()

        # Coerce potentially complex signature objects into JSON/string for SQLite binding
        sig_value = getattr(event, "signature", None)
        if isinstance(sig_value, (dict, list)):
            try:
                sig_value = json.dumps(sig_value)
            except Exception:
                sig_value = str(sig_value)

        # Extract receipt/alert fields if present
        verdict = getattr(event, 'verdict', None)
        cost = json.dumps(getattr(event, 'cost', {})) if hasattr(event, 'cost') else None
        latency_ms = getattr(event, 'latency_ms', None)
        refusal_reason = getattr(event, 'refusal_reason', None)
        severity = getattr(event, 'severity', None)
        alert_type = getattr(event, 'alert_type', None)
        conditions = json.dumps(getattr(event, 'conditions', {})) if hasattr(event, 'conditions') else None
        actions = json.dumps(getattr(event, 'recommended_actions', [])) if hasattr(event, 'recommended_actions') else None

        self.sqlite_conn.execute(
            """INSERT OR REPLACE INTO events (
                event_id, event_type, tenant_id, valid_time, transaction_time,
                observed_at, dip_hash, signature, parent_ids, namespaces,
                data, privacy_tags, source, tier, chi_score, dignity_score,
                traceparent, tracestate, verdict, cost, latency_ms,
                refusal_reason, severity, alert_type, conditions, recommended_actions
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                event.event_id,
                event.event_type,
                event.tenant_id,
                valid_ts,
                trans_ts,
                obs_ts,
                event.dip_hash,
                sig_value,
                parent_ids_json,
                namespaces_json,
                data_json,
                privacy_tags_json,
                event.source,
                event.tier.name if hasattr(event, "tier") else AlertTier.RECORDED.name,
                float(getattr(event, "chi_score", 0.0)),
                getattr(event, "dignity_score", None),
                getattr(event, "traceparent", None),
                getattr(event, "tracestate", None),
                verdict,
                cost,
                latency_ms,
                refusal_reason,
                severity,
                alert_type,
                conditions,
                actions,
            )
        )

    def _buffer_cold(self, event: BaseEvent) -> None:
        """Append event to Parquet buffer respecting batch size."""
        self._write_buffer.append(event.to_parquet_row())
        if len(self._write_buffer) >= self.batch_size:
            self.flush_parquet()

    def flush_parquet(self):
        """Flush write buffer to Parquet files."""
        if not self._write_buffer:
            return

        if not HAS_PARQUET:
            logger.warning("PyArrow not available, skipping Parquet write")
            self._write_buffer.clear()
            return

        # Group by partition (year/month/day/tenant/type)
        partitions: Dict[tuple, List[Dict]] = {}

        for row in self._write_buffer:
            key = (
                row['year'],
                row['month'],
                row['day'],
                row.get('tenant_id', 'default'),
                row.get('event_type', 'unknown').split('.')[-2] if '.' in row.get('event_type', '') else 'unknown'
            )
            if key not in partitions:
                partitions[key] = []
            partitions[key].append(row)

        # Write each partition
        for (year, month, day, tenant, etype), rows in partitions.items():
            partition_dir = (
                self.cold_dir /
                f"year={year}" /
                f"month={month:02d}" /
                f"day={day:02d}" /
                f"tenant={tenant}" /
                f"type={etype}"
            )
            partition_dir.mkdir(parents=True, exist_ok=True)

            # Generate filename with timestamp
            timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
            filepath = partition_dir / f"events_{timestamp}.parquet"

            # Convert to PyArrow table
            table = pa.Table.from_pylist(rows)

            # Write Parquet with compression
            pq.write_table(
                table,
                filepath,
                compression='zstd',
                compression_level=3,
            )

            logger.info(f"Wrote {len(rows)} events to {filepath}")

        # Clear buffer
        self._write_buffer.clear()

    def query_recent(
        self,
        tenant: str,
        hours: int = 24,
        event_type: Optional[str] = None,
        namespaces: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Query recent events from hot tier (SQLite).

        Args:
            tenant: Tenant ID
            hours: How many hours back to query
            event_type: Optional event type filter
            namespaces: Optional namespace filter

        Returns:
            List of event dicts
        """
        cutoff = datetime.now(timezone.utc) - timedelta(hours=hours)
        cutoff_ts = cutoff.timestamp()

        query = "SELECT * FROM events WHERE tenant_id = ? AND valid_time >= ?"
        params = [tenant, cutoff_ts]

        if event_type:
            query += " AND event_type = ?"
            params.append(event_type)

        if namespaces:
            # SQLite JSON containment check
            for ns in namespaces:
                query += " AND namespaces LIKE ?"
                params.append(f'%"{ns}"%')

        query += " ORDER BY valid_time DESC"

        cursor = self.sqlite_conn.execute(query, params)
        rows = cursor.fetchall()

        # Convert to dicts
        columns = [desc[0] for desc in cursor.description]
        return [dict(zip(columns, row)) for row in rows]

    def query_analytical(self, sql: str) -> Any:
        """
        Query cold tier (Parquet) using DuckDB.

        Args:
            sql: DuckDB SQL query

        Returns:
            DuckDB result (can be converted to pandas/arrow/dict)

        Example:
            df = store.query_analytical('''
                SELECT event_type, COUNT(*) as count
                FROM read_parquet('data/cold/events/**/*.parquet')
                WHERE year = 2025 AND month = 11
                GROUP BY event_type
            ''')
        """
        if not self.duckdb_conn:
            raise RuntimeError("DuckDB not available")

        return self.duckdb_conn.execute(sql)

    def query_evidence_chain(self, event_id: str, max_depth: int = 10) -> List[Dict[str, Any]]:
        """
        Trace evidence chain backwards from event.

        Follows parent_ids to build full provenance DAG.
        """
        chain = []
        visited = set()
        queue = [event_id]

        while queue and len(chain) < max_depth:
            current_id = queue.pop(0)

            if current_id in visited:
                continue

            visited.add(current_id)

            # Query from hot tier first
            cursor = self.sqlite_conn.execute(
                "SELECT * FROM events WHERE event_id = ?",
                (current_id,)
            )
            row = cursor.fetchone()

            if row:
                columns = [desc[0] for desc in cursor.description]
                event_dict = dict(zip(columns, row))
                chain.append(event_dict)

                # Parse parent IDs
                parent_ids_json = event_dict.get('parent_ids', '[]')
                try:
                    parent_ids = json.loads(parent_ids_json)
                    queue.extend(parent_ids)
                except json.JSONDecodeError:
                    pass

        return chain

    def prune_hot_tier(self):
        """
        Remove old events from SQLite (hot tier).

        Keeps last N days based on hot_retention_days.
        Cold tier (Parquet) retains everything.
        """
        cutoff = datetime.now(timezone.utc) - timedelta(days=self.hot_retention_days)
        cutoff_ts = cutoff.timestamp()

        cursor = self.sqlite_conn.execute(
            "DELETE FROM events WHERE valid_time < ?",
            (cutoff_ts,)
        )

        deleted = cursor.rowcount
        logger.info(f"Pruned {deleted} events from hot tier (older than {self.hot_retention_days} days)")

        # Vacuum to reclaim space
        self.sqlite_conn.execute("VACUUM")

    def close(self):
        """Close connections and flush buffers."""
        self.flush_parquet()
        self.sqlite_conn.close()
        if self.duckdb_conn:
            self.duckdb_conn.close()


# === CLI for testing ===

if __name__ == "__main__":
    import argparse
    from receipts.base_event import BaseEvent, AlertEvent

    parser = argparse.ArgumentParser(description="Event Store CLI")
    parser.add_argument("--data-dir", type=Path, default=Path("data"), help="Data directory")
    parser.add_argument("--write-test", action="store_true", help="Write test events")
    parser.add_argument("--query-recent", type=int, help="Query recent N hours")
    parser.add_argument("--prune", action="store_true", help="Prune old events")

    args = parser.parse_args()

    store = EventStore(data_dir=args.data_dir)

    if args.write_test:
        print("Writing test events...")

        # Test event 1: Clinical prophecy
        event1 = BaseEvent(
            event_type="loom.prophecy.clinical.sepsis.v1",
            tenant_id="test_hospital",
            namespaces=["icu"],
            data={"prediction": 0.85, "patient_id": "PT123"},
        )
        store.write(event1)

        # Test event 2: Alert
        event2 = AlertEvent(
            event_type="loom.alert.drift.v1",
            tenant_id="test_hospital",
            namespaces=["icu"],
            severity="warning",
            alert_type="calibration_drift",
            data={"ece_delta": 0.07},
        )
        store.write(event2)

        store.flush_parquet()
        print(f"Wrote 2 test events (event IDs: {event1.event_id}, {event2.event_id})")

    if args.query_recent:
        print(f"\nQuerying last {args.query_recent} hours...")
        events = store.query_recent(tenant="test_hospital", hours=args.query_recent)
        print(f"Found {len(events)} events:")
        for event in events[:10]:  # Show first 10
            print(f"  - {event['event_id']}: {event['event_type']}")

    if args.prune:
        print("\nPruning old events...")
        store.prune_hot_tier()

    store.close()
    print("\nDone!")
