#!/usr/bin/env python3
# INFRA:DEAD rds -- AWS RDS is gone. Any file referencing rds_connect.py or this hostname is stale and must be ported.
"""
rrc_bosonic_db_buffer.py — Asynchronous buffering database ingestion program.

Queues RRC simulation receipts and writes them in buffered batches to the PostgreSQL database.
Uses asyncio for scheduling and delegates synchronous database operations to a thread pool.
"""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import logging
import os
import sys
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Enable importing sibling modules
sys.path.insert(0, str(Path(__file__).parent))

try:
    import psycopg2
    import psycopg2.extras
    from rds_connect import connect_rds
except ImportError:
    # Fallback/Mock for local testing when psycopg2 is missing
    psycopg2 = None
    connect_rds = None

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
)
log = logging.getLogger("db_buffer")


def get_insert_query(table: str, keys: List[str]) -> str:
    """Generates parameterized INSERT queries with appropriate conflict resolution."""
    fields = ", ".join(keys)
    placeholders = ", ".join([f"%({k})s" for k in keys])

    if table == "ene.packages":
        return f"""
            INSERT INTO {table} ({fields})
            VALUES ({placeholders})
            ON CONFLICT (pkg) DO UPDATE SET
                title = EXCLUDED.title,
                content = EXCLUDED.content,
                content_hash = EXCLUDED.content_hash,
                updated_at = EXCLUDED.updated_at,
                verification_status = EXCLUDED.verification_status,
                promotion_state = EXCLUDED.promotion_state
        """
    elif table == "ene.receipts":
        return f"""
            INSERT INTO {table} ({fields})
            VALUES ({placeholders})
            ON CONFLICT (id) DO NOTHING
        """
    elif table == "ene.crossing_weights":
        return f"""
            INSERT INTO {table} ({fields})
            VALUES ({placeholders})
            ON CONFLICT (recipe_id, row_index, col_index) DO UPDATE SET
                weight_raw = EXCLUDED.weight_raw,
                contractive_factor = EXCLUDED.contractive_factor
        """
    elif table == "ene.eigensolid_snapshots":
        return f"""
            INSERT INTO {table} ({fields})
            VALUES ({placeholders})
            ON CONFLICT (id) DO NOTHING
        """
    elif table == "ene.braid_strands":
        return f"""
            INSERT INTO {table} ({fields})
            VALUES ({placeholders})
            ON CONFLICT (snapshot_id, strand_index) DO UPDATE SET
                phase_x = EXCLUDED.phase_x,
                phase_y = EXCLUDED.phase_y
        """
    else:
        # Default fallback
        return f"INSERT INTO {table} ({fields}) VALUES ({placeholders}) ON CONFLICT DO NOTHING"


class AsyncDatabaseBuffer:
    """Asynchronously buffers and writes database records in thread-safe batches."""

    def __init__(
        self,
        batch_size: int = 10,
        flush_interval: float = 2.0,
        max_retries: int = 5,
        retry_delay_base: float = 1.0,
        dry_run: bool = False,
    ):
        self.batch_size = batch_size
        self.flush_interval = flush_interval
        self.max_retries = max_retries
        self.retry_delay_base = retry_delay_base
        self.dry_run = dry_run

        self.queue: asyncio.Queue[Tuple[str, Dict[str, Any]]] = asyncio.Queue()
        self.conn: Optional[Any] = None
        self.worker_task: Optional[asyncio.Task] = None
        self.running = False
        self._last_flush_time = time.time()

    async def start(self):
        """Starts the background worker task."""
        self.running = True
        self.worker_task = asyncio.create_task(self._worker_loop())
        log.info("Async database buffer coordinator started.")

    async def stop(self):
        """Stops the worker and flushes any remaining items."""
        self.running = False
        if self.worker_task:
            # Let the worker finish processing what's in the queue
            self.worker_task.cancel()
            try:
                await self.worker_task
            except asyncio.CancelledError:
                pass

        # Flush anything left over
        await self._flush_remaining()

        # Close connection
        if self.conn:
            try:
                await asyncio.to_thread(self.conn.close)
                log.info("Closed database connection.")
            except Exception as e:
                log.error(f"Error closing database connection: {e}")
            self.conn = None

    async def enqueue(self, table_name: str, row_data: Dict[str, Any]):
        """Puts an item to be inserted into the queue."""
        await self.queue.put((table_name, row_data))

    async def _worker_loop(self):
        """Main loop that gathers items and triggers flushes."""
        batch: List[Tuple[str, Dict[str, Any]]] = []

        while self.running:
            try:
                # Wait for an item, or timeout to trigger periodic flush
                timeout = max(0.1, self.flush_interval - (time.time() - self._last_flush_time))
                try:
                    item = await asyncio.wait_for(self.queue.get(), timeout=timeout)
                    batch.append(item)
                    self.queue.task_done()
                except asyncio.TimeoutError:
                    # Flush interval reached
                    pass

                # Flush conditions
                if len(batch) >= self.batch_size or (batch and (time.time() - self._last_flush_time >= self.flush_interval)):
                    await self._flush_batch_with_retry(batch)
                    batch = []
                    self._last_flush_time = time.time()

            except asyncio.CancelledError:
                break
            except Exception as e:
                log.error(f"Error in worker loop: {e}", exc_info=True)
                await asyncio.sleep(1.0)

        # Catch remaining items
        if batch:
            await self._flush_batch_with_retry(batch)

    async def _flush_remaining(self):
        """Drains the queue and flushes all items."""
        batch = []
        while not self.queue.empty():
            item = self.queue.get_nowait()
            batch.append(item)
            self.queue.task_done()

        if batch:
            await self._flush_batch_with_retry(batch)

    async def _flush_batch_with_retry(self, batch: List[Tuple[str, Dict[str, Any]]]):
        """Retries batch flushing with exponential backoff on database failure."""
        if not batch:
            return

        log.info(f"Flushing batch of {len(batch)} records...")
        if self.dry_run:
            log.info(f"[DRY-RUN] Would have written {len(batch)} records.")
            for table, data in batch:
                log.debug(f"[DRY-RUN] Table: {table}, Keys: {list(data.keys())}")
            return

        for attempt in range(1, self.max_retries + 1):
            try:
                await asyncio.to_thread(self._flush_batch_sync, batch)
                log.info("Batch flush successful.")
                return
            except Exception as e:
                delay = self.retry_delay_base * (2 ** (attempt - 1))
                log.warning(f"Database write attempt {attempt}/{self.max_retries} failed: {e}. Retrying in {delay:.1f}s...")
                # Reset connection on failure so it attempts to reconnect
                if self.conn:
                    try:
                        self.conn.close()
                    except Exception:
                        pass
                    self.conn = None
                await asyncio.sleep(delay)

        log.error(f"Failed to write batch of {len(batch)} records after {self.max_retries} attempts.")
        # Optionally dump to a local fallback log file so we don't lose data
        fallback_file = Path("shared-data/db_buffer_fallback.jsonl")
        try:
            fallback_file.parent.mkdir(parents=True, exist_ok=True)
            with open(fallback_file, "a") as f:
                for table, data in batch:
                    f.write(json.dumps({"table": table, "data": data, "failed_at": datetime.now(timezone.utc).isoformat()}) + "\n")
            log.error(f"Dumped failed batch to fallback log: {fallback_file}")
        except Exception as fe:
            log.critical(f"Could not write to fallback log: {fe}")

    def _flush_batch_sync(self, batch: List[Tuple[str, Dict[str, Any]]]):
        """Synchronous flush implementation. Executed inside a thread pool."""
        if psycopg2 is None or connect_rds is None:
            raise RuntimeError("psycopg2 is not installed or rds_connect is missing.")

        if not self.conn:
            log.info("Establishing database connection...")
            self.conn = connect_rds(connect_timeout=5)

        cur = self.conn.cursor()
        try:
            # Group entries by table name to optimize execution
            grouped: Dict[str, List[Dict[str, Any]]] = {}
            for table, data in batch:
                grouped.setdefault(table, []).append(data)

            for table, rows in grouped.items():
                if not rows:
                    continue
                # Extract keys from the first row (assuming all rows in the group have identical keys)
                keys = list(rows[0].keys())
                query = get_insert_query(table, keys)

                log.debug(f"Executing batch insert on {table} for {len(rows)} rows.")
                psycopg2.extras.execute_batch(cur, query, rows)

            self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            # Catch duplicate content hash violations and skip them cleanly as successful dedups
            if e.__class__.__name__ == "UniqueViolation" and "ene_pkg_hash_idx" in str(e):
                log.info("Deduplication filter triggered: record with identical content_hash already exists. Skipping.")
                return
            raise e
        finally:
            cur.close()


# ── Helper to convert simulation receipts into database rows ─────────────────

def receipt_to_db_rows(receipt: Dict[str, Any]) -> List[Tuple[str, Dict[str, Any]]]:
    """Converts a standard RRC simulation receipt into ene.packages and ene.receipts rows."""
    rows = []

    # 1. Package Row
    pkg_id = f"pkg_bosonic_tn_N{receipt['hardware_details']['N']}_{str(uuid.uuid4())[:8]}"
    content = json.dumps(receipt, sort_keys=True)
    content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()

    pkg_row = {
        "pkg": pkg_id,
        "package_type": "bosonic_tensor_network",
        "title": f"Bosonic Tensor Network Centrality (N = {receipt['hardware_details']['N']})",
        "content": content,
        "content_hash": content_hash,
        "element": "METAL",
        "tags": json.dumps(["bosonic_tensor", "centrality", f"N_{receipt['hardware_details']['N']}"]),
        "source": "rrc_bosonic_tensor_gpu.py",
        "created_at": receipt.get("computed_at", datetime.now(timezone.utc).isoformat()),
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "verification_status": "verified",
        "promotion_state": "held",
        "domain": "RRC",
        "archetype": "receipt",
    }
    rows.append(("ene.packages", pkg_row))

    # 2. Receipt Row
    rec_id = f"rec_bosonic_tn_N{receipt['hardware_details']['N']}_{str(uuid.uuid4())[:8]}"

    # Derive input hash (from hardware details) and output hash (from receipt sha256)
    input_preimage = json.dumps(receipt["hardware_details"], sort_keys=True)
    input_hash = hashlib.sha256(input_preimage.encode("utf-8")).hexdigest()

    rec_row = {
        "id": rec_id,
        "package_id": pkg_id,
        "receipt_type": receipt.get("schema", "rrc_bosonic_tensor_gpu_receipt_v1"),
        "receipt_hash": receipt.get("receipt_sha256", content_hash),
        "input_hash": input_hash,
        "output_hash": receipt.get("receipt_sha256", content_hash),
        "toolchain": "wgpu-vulkan-rk4",
        "verifier": "RTX 4070 SUPER",
        "status": "verified",
        "residual": json.dumps(receipt.get("entropies", [])),
        "created_at": receipt.get("computed_at", datetime.now(timezone.utc).isoformat()),
    }
    rows.append(("ene.receipts", rec_row))

    return rows


# ── CLI & Main Entry ──────────────────────────────────────────────────────────

async def async_main(args) -> int:
    # Initialize the buffer coordinator
    buffer = AsyncDatabaseBuffer(
        batch_size=args.batch_size,
        flush_interval=args.flush_interval,
        max_retries=args.max_retries,
        dry_run=args.dry_run,
    )
    await buffer.start()

    try:
        # Check if we are reading a specific receipt file
        if args.file:
            path = Path(args.file)
            if not path.exists():
                log.error(f"Input file does not exist: {path}")
                return 1

            with open(path, "r") as f:
                receipt = json.load(f)

            rows = receipt_to_db_rows(receipt)
            for table, row in rows:
                await buffer.enqueue(table, row)
            log.info(f"Queued receipt from file: {path}")

        # Check if we are reading from stdin (interactive/piped stream)
        elif args.stdin:
            log.info("Reading JSON/JSONL records from standard input (Ctrl+D to finish)...")
            loop = asyncio.get_running_loop()

            # Read line-by-line asynchronously
            reader = asyncio.StreamReader()
            protocol = asyncio.StreamReaderProtocol(reader)
            await loop.connect_read_pipe(lambda: protocol, sys.stdin)

            while True:
                line_bytes = await reader.readline()
                if not line_bytes:
                    break
                line = line_bytes.decode("utf-8").strip()
                if not line:
                    continue

                try:
                    payload = json.loads(line)
                    # Support both raw table rows or formatted receipts
                    if "schema" in payload and "hardware_details" in payload:
                        rows = receipt_to_db_rows(payload)
                        for table, row in rows:
                            await buffer.enqueue(table, row)
                    elif "table" in payload and "data" in payload:
                        await buffer.enqueue(payload["table"], payload["data"])
                    else:
                        # Direct dictionary input -> treat as a general package
                        pkg_id = f"pkg_stream_{str(uuid.uuid4())[:8]}"
                        pkg_row = {
                            "pkg": pkg_id,
                            "package_type": "stream_item",
                            "title": "Streamed Database Payload",
                            "content": json.dumps(payload),
                            "content_hash": hashlib.sha256(json.dumps(payload).encode("utf-8")).hexdigest(),
                            "element": "GROUND",
                            "created_at": datetime.now(timezone.utc).isoformat(),
                            "updated_at": datetime.now(timezone.utc).isoformat(),
                        }
                        await buffer.enqueue("ene.packages", pkg_row)
                except Exception as e:
                    log.error(f"Failed to parse stdin payload: {e}")

    finally:
        # Stop will flush all queued items
        log.info("Stopping coordinator and flushing remaining items...")
        await buffer.stop()

    return 0


def main():
    parser = argparse.ArgumentParser(description="Async Database Ingestion Buffer Shim")
    parser.add_argument("--batch-size", type=int, default=10, help="Batch size for writing to DB (default: 10)")
    parser.add_argument("--flush-interval", type=float, default=2.0, help="Interval in seconds for periodic flush (default: 2.0)")
    parser.add_argument("--max-retries", type=int, default=5, help="Max database reconnection/write retries (default: 5)")
    parser.add_argument("--dry-run", action="store_true", help="Run without writing to database (dry run)")
    parser.add_argument("--file", type=str, help="Ingest a specific JSON receipt file")
    parser.add_argument("--stdin", action="store_true", help="Stream payloads line-by-line from standard input")

    args = parser.parse_args()

    # If neither file nor stdin is provided, default to stdin if not a tty, otherwise print help
    if not args.file and not args.stdin:
        if not sys.stdin.isatty():
            args.stdin = True
        else:
            parser.print_help()
            sys.exit(1)

    try:
        sys.exit(asyncio.run(async_main(args)))
    except KeyboardInterrupt:
        log.info("Buffer program terminated by user.")
        sys.exit(130)


if __name__ == "__main__":
    main()