local_pipeline.py

"""Local pipeline: end-to-end data flow without Kafka or Flink.

Generates → validates → enriches → writes to DuckDB in real-time.
Proves the pipeline works end-to-end, locally, with zero infrastructure.

Usage:
    python -m src.processing.local_pipeline                # default: 10 events/sec
    python -m src.processing.local_pipeline --eps 50       # 50 events/sec
    python -m src.processing.local_pipeline --burst 500    # one-shot: 500 events
"""

import argparse
import json
import os
import time
from datetime import UTC, datetime
from pathlib import Path

import duckdb
import structlog
import yaml
from pyiceberg.exceptions import NoSuchPropertyException, RESTError, ValidationError

from src.ingestion.producers.event_producer import (
    generate_click,
    generate_order,
    generate_payment,
    generate_product,
)
from src.logger import configure_logging
from src.processing.iceberg_sink import IcebergSink
from src.processing.transformations.enrichment import (
    compute_payment_risk_score,
    enrich_clickstream,
    enrich_order,
)
from src.quality.validators.schema_validator import validate_event
from src.quality.validators.semantic_validator import validate_semantics

DB_PATH = os.getenv("DUCKDB_PATH", "agentflow_demo.duckdb")


def _ensure_tables(conn: duckdb.DuckDBPyConnection) -> None:
    """Create all tables if they don't exist."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS orders_v2 (
            order_id VARCHAR PRIMARY KEY,
            user_id VARCHAR,
            status VARCHAR,
            total_amount DECIMAL(10,2),
            currency VARCHAR DEFAULT 'USD',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS products_current (
            product_id VARCHAR PRIMARY KEY,
            name VARCHAR,
            category VARCHAR,
            price DECIMAL(10,2),
            in_stock BOOLEAN DEFAULT TRUE,
            stock_quantity INTEGER DEFAULT 0
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS sessions_aggregated (
            session_id VARCHAR PRIMARY KEY,
            user_id VARCHAR,
            started_at TIMESTAMP,
            ended_at TIMESTAMP,
            duration_seconds FLOAT,
            event_count INTEGER,
            unique_pages INTEGER,
            funnel_stage VARCHAR,
            is_conversion BOOLEAN DEFAULT FALSE
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS users_enriched (
            user_id VARCHAR PRIMARY KEY,
            total_orders INTEGER DEFAULT 0,
            total_spent DECIMAL(10,2) DEFAULT 0,
            first_order_at TIMESTAMP,
            last_order_at TIMESTAMP,
            preferred_category VARCHAR
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS pipeline_events (
            event_id VARCHAR,
            topic VARCHAR,
            tenant_id VARCHAR DEFAULT 'default',
            event_type VARCHAR,
            latency_ms INTEGER,
            processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.execute(
        "ALTER TABLE pipeline_events ADD COLUMN IF NOT EXISTS tenant_id VARCHAR DEFAULT 'default'"
    )


def _event_tenant(event: dict) -> str:
    source_metadata = event.get("source_metadata", {})
    metadata_tenant = source_metadata.get("tenant") if isinstance(source_metadata, dict) else None
    tenant = event.get("tenant") or metadata_tenant
    return str(tenant) if tenant else "default"


def _process_event(
    conn: duckdb.DuckDBPyConnection,
    event: dict,
    iceberg_sink: IcebergSink | None = None,
) -> tuple[bool, str]:
    """Validate, enrich, and store a single event. Returns (success, reason)."""
    event_type = event.get("event_type", "")
    event_id = event.get("event_id", "unknown")
    tenant_id = _event_tenant(event)

    conn.execute("BEGIN")
    try:
        # Schema validation
        schema_result = validate_event(event)
        if not schema_result.is_valid:
            conn.execute(
                """
                INSERT INTO pipeline_events (
                    event_id, topic, tenant_id, event_type, latency_ms, processed_at
                )
                VALUES (?, 'events.deadletter', ?, ?, 0, ?)
                """,
                [event_id, tenant_id, event_type, datetime.now(UTC)],
            )
            if iceberg_sink is not None:
                iceberg_sink.write_batch(
                    "dead_letter",
                    [
                        {
                            "event_id": event.get("event_id"),
                            "event_type": event.get("event_type"),
                            "reason": f"schema: {schema_result.errors[0]}",
                            "source_topic": "events.deadletter",
                            "received_at": datetime.now(UTC),
                            "payload": event,
                        }
                    ],
                )
            conn.execute("COMMIT")
            return False, f"schema: {schema_result.errors[0]}"

        # Semantic validation
        semantic_result = validate_semantics(event)
        error_issues = [i for i in semantic_result.issues if i.severity == "error"]
        if error_issues:
            conn.execute(
                """
                INSERT INTO pipeline_events (
                    event_id, topic, tenant_id, event_type, latency_ms, processed_at
                )
                VALUES (?, 'events.deadletter', ?, ?, 0, ?)
                """,
                [event_id, tenant_id, event_type, datetime.now(UTC)],
            )
            if iceberg_sink is not None:
                iceberg_sink.write_batch(
                    "dead_letter",
                    [
                        {
                            "event_id": event.get("event_id"),
                            "event_type": event.get("event_type"),
                            "reason": f"semantic: {error_issues[0].rule}",
                            "source_topic": "events.deadletter",
                            "received_at": datetime.now(UTC),
                            "payload": event,
                        }
                    ],
                )
            conn.execute("COMMIT")
            return False, f"semantic: {error_issues[0].rule}"

        # Enrichment
        if event_type.startswith("order."):
            event = enrich_order(event)
            _upsert_order(conn, event)
            if iceberg_sink is not None:
                iceberg_sink.write_batch("orders", [event])
        elif event_type in ("click", "page_view", "add_to_cart"):
            event = enrich_clickstream(event)
            _upsert_session(conn, event)
            if iceberg_sink is not None:
                iceberg_sink.write_batch("clickstream", [event])
        elif event_type.startswith("payment."):
            event = compute_payment_risk_score(event)
            if iceberg_sink is not None:
                iceberg_sink.write_batch("payments", [event])
        elif event_type.startswith("product."):
            _upsert_product(conn, event)
            if iceberg_sink is not None:
                iceberg_sink.write_batch("inventory", [event])

        # Record in pipeline_events
        ts = event.get("timestamp", "")
        try:
            event_ts = datetime.fromisoformat(ts)
            if event_ts.tzinfo is None:
                event_ts = event_ts.replace(tzinfo=UTC)
            latency_ms = int((datetime.now(UTC) - event_ts).total_seconds() * 1000)
        except (ValueError, TypeError):
            latency_ms = 0

        conn.execute(
            """
            INSERT INTO pipeline_events (
                event_id, topic, tenant_id, event_type, latency_ms, processed_at
            )
            VALUES (?, 'events.validated', ?, ?, ?, ?)
            """,
            [event_id, tenant_id, event_type, latency_ms, datetime.now(UTC)],
        )
        conn.execute("COMMIT")
        return True, "ok"
    # rollback must preserve the original pipeline failure
    except Exception:  # nosec B110
        # Transaction rollback must happen before unexpected errors propagate.
        conn.execute("ROLLBACK")
        raise


def _upsert_order(conn: duckdb.DuckDBPyConnection, event: dict) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO orders_v2
        (order_id, user_id, status, total_amount, currency, created_at)
        VALUES (?, ?, ?, ?, ?, ?)
    """,
        [
            event["order_id"],
            event["user_id"],
            event["status"],
            float(event["total_amount"]),
            event.get("currency", "USD"),
            datetime.fromisoformat(event["timestamp"]),
        ],
    )
    # Update user aggregate
    conn.execute(
        """
        INSERT OR REPLACE INTO users_enriched
        (user_id, total_orders, total_spent,
         first_order_at, last_order_at, preferred_category)
        SELECT
            user_id,
            COUNT(*) as total_orders,
            SUM(total_amount) as total_spent,
            MIN(created_at),
            MAX(created_at),
            NULL
        FROM orders_v2
        WHERE user_id = ? AND status != 'cancelled'
        GROUP BY user_id
    """,
        [event["user_id"]],
    )


def _upsert_product(conn: duckdb.DuckDBPyConnection, event: dict) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO products_current
        (product_id, name, category, price, in_stock, stock_quantity)
        VALUES (?, ?, ?, ?, ?, ?)
    """,
        [
            event["product_id"],
            event["name"],
            event["category"],
            float(event["price"]),
            event["in_stock"],
            event["stock_quantity"],
        ],
    )


def _upsert_session(conn: duckdb.DuckDBPyConnection, event: dict) -> None:
    session_id = event.get("session_id", "unknown")
    derived = event.get("_derived", {})
    page_cat = derived.get("page_category", "other")

    # Determine funnel stage from page category
    stage_order = {
        "checkout": 4,
        "cart": 3,
        "product_detail": 2,
        "search": 1,
        "home": 0,
        "other": 0,
    }
    new_stage_val = stage_order.get(page_cat, 0)

    existing = conn.execute(
        "SELECT funnel_stage, event_count FROM sessions_aggregated WHERE session_id = ?",
        [session_id],
    ).fetchone()

    if existing:
        old_stage = existing[0] or "bounce"
        old_count = existing[1] or 0
        old_stage_val = stage_order.get(old_stage, 0)
        funnel = page_cat if new_stage_val > old_stage_val else old_stage
        conn.execute(
            """
            UPDATE sessions_aggregated
            SET event_count = ?,
                funnel_stage = ?,
                is_conversion = ?
            WHERE session_id = ?
        """,
            [
                old_count + 1,
                funnel,
                funnel == "checkout",
                session_id,
            ],
        )
    else:
        conn.execute(
            """
            INSERT INTO sessions_aggregated
            (session_id, user_id, started_at, ended_at,
             duration_seconds, event_count, unique_pages,
             funnel_stage, is_conversion)
            VALUES (?, ?, ?, NULL, 0, 1, 1, ?, ?)
        """,
            [
                session_id,
                event.get("user_id"),
                datetime.now(UTC),
                page_cat,
                page_cat == "checkout",
            ],
        )


def _generate_random_event() -> tuple[str, dict]:
    """Generate a random event using existing producers."""
    import random

    generators: list[tuple] = [
        (0.15, generate_order),
        (0.25, generate_payment),
        (0.95, generate_click),
        (1.00, generate_product),
    ]
    roll = random.random()
    for threshold, gen in generators:
        if roll < threshold:
            topic, event = gen()
            return topic, json.loads(event.model_dump_json())
    topic, event = generate_product()
    return topic, json.loads(event.model_dump_json())


def _format_rate(total: float, elapsed: float) -> str:
    # Guard against a zero elapsed window: on coarse-resolution monotonic clocks
    # (e.g. Windows) the first 100-event progress tick of a --burst run can land
    # within a single clock tick (elapsed == 0.0), which would raise
    # ZeroDivisionError in the progress log. (audit_28_06_26.md §5 low)
    return f"{total / max(elapsed, 0.001):.0f} evt/s"


def run(events_per_second: int = 10, burst: int = 0) -> None:
    """Run the local pipeline."""
    configure_logging()
    logger = structlog.get_logger()
    conn = duckdb.connect(DB_PATH)
    _ensure_tables(conn)
    iceberg_sink = None
    iceberg_config = os.getenv("AGENTFLOW_ICEBERG_CONFIG")
    if not iceberg_config:
        default_iceberg_config = Path("config/iceberg.yaml")
        if default_iceberg_config.exists():
            iceberg_config = str(default_iceberg_config)
    if iceberg_config:
        try:
            iceberg_sink = IcebergSink(config_path=iceberg_config)
            iceberg_sink.create_tables_if_not_exist()
        except (
            OSError,
            KeyError,
            ValueError,
            yaml.YAMLError,
            NoSuchPropertyException,
            RESTError,
            ValidationError,
        ) as exc:
            iceberg_sink = None
            logger.warning(
                "iceberg_sink_unavailable",
                config=iceberg_config,
                error=str(exc),
                exc_info=True,
            )

    logger.info(
        "local_pipeline_started",
        db=DB_PATH,
        eps=events_per_second,
        burst=burst,
    )

    total = 0
    valid = 0
    invalid = 0
    start_time = time.monotonic()

    try:
        count = burst if burst > 0 else float("inf")
        while total < count:
            _, event = _generate_random_event()
            success, reason = _process_event(conn, event, iceberg_sink=iceberg_sink)

            total += 1
            if success:
                valid += 1
            else:
                invalid += 1

            if total % 100 == 0:
                elapsed = time.monotonic() - start_time
                logger.info(
                    "pipeline_progress",
                    total=total,
                    valid=valid,
                    invalid=invalid,
                    rate=_format_rate(total, elapsed),
                )

            if burst == 0:
                time.sleep(1.0 / events_per_second)

    except KeyboardInterrupt:
        pass
    finally:
        elapsed = time.monotonic() - start_time
        conn.close()
        logger.info(
            "local_pipeline_stopped",
            total=total,
            valid=valid,
            invalid=invalid,
            duration_s=round(elapsed, 1),
            avg_rate=_format_rate(total, elapsed),
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AgentFlow local pipeline")
    parser.add_argument("--eps", type=int, default=10, help="Events per second")
    parser.add_argument("--burst", type=int, default=0, help="One-shot: N events then stop")
    args = parser.parse_args()
    run(events_per_second=args.eps, burst=args.burst)