"""
SessionLedger — VAULT999 Postgres Persistence Layer

Provides Merkle-chained audit trail for constitutional verdicts.
Schema: VAULT999 v3 (hybrid)

DITEMPA BUKAN DIBERI
"""

import os
import json
import hashlib
from datetime import datetime, timezone
from typing import Optional, Dict, Any, List
from dataclasses import dataclass

# Try to import asyncpg for Postgres support
try:
    import asyncpg

    ASYNCPG_AVAILABLE = True
except ImportError:
    ASYNCPG_AVAILABLE = False


@dataclass
class VaultEntry:
    """VAULT999 v3 entry schema."""

    entry_id: str
    session_id: str
    timestamp: str
    verdict_type: str  # SEAL | VOID | PARTIAL | SABAR | 888_HOLD
    risk_level: str  # LOW | MEDIUM | HIGH | CRITICAL
    category: str
    environment: str
    payload: Dict[str, Any]
    query_summary: str
    prev_hash: str
    entry_hash: str
    schema_version: str = "3.0"
    floors_checked: List[str] = None
    floors_failed: List[str] = None

    def __post_init__(self):
        if self.floors_checked is None:
            self.floors_checked = []
        if self.floors_failed is None:
            self.floors_failed = []


# Compatibility alias for legacy code
SessionEntry = VaultEntry


class SessionLedger:
    """
    VAULT999 Ledger — Immutable Merkle-chained audit trail.

    Supports:
    - Postgres backend (production)
    - In-memory fallback (development/testing)
    """

    # SQL for creating the vault table
    CREATE_TABLE_SQL = """
    CREATE TABLE IF NOT EXISTS vault999 (
        id SERIAL PRIMARY KEY,
        entry_id VARCHAR(64) UNIQUE NOT NULL,
        session_id VARCHAR(128) NOT NULL,
        timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),
        verdict_type VARCHAR(20) NOT NULL,
        risk_level VARCHAR(20) NOT NULL,
        category VARCHAR(64),
        environment VARCHAR(20),
        payload JSONB,
        query_summary TEXT,
        prev_hash VARCHAR(64),
        entry_hash VARCHAR(64) NOT NULL,
        schema_version VARCHAR(10) DEFAULT '3.0',
        floors_checked TEXT[],
        floors_failed TEXT[],
        created_at TIMESTAMPTZ DEFAULT NOW()
    );
    
    CREATE INDEX IF NOT EXISTS idx_vault999_session ON vault999(session_id);
    CREATE INDEX IF NOT EXISTS idx_vault999_verdict ON vault999(verdict_type);
    CREATE INDEX IF NOT EXISTS idx_vault999_timestamp ON vault999(timestamp DESC);
    """

    def __init__(self, database_url: Optional[str] = None):
        self.database_url = database_url or os.environ.get("DATABASE_URL")
        self._pool: Optional[asyncpg.Pool] = None
        self._memory_ledger: List[VaultEntry] = []
        self._last_hash = "GENESIS"
        self._initialized = False

    @property
    def is_postgres_available(self) -> bool:
        """Check if Postgres is configured and available."""
        return ASYNCPG_AVAILABLE and bool(self.database_url)

    async def initialize(self) -> bool:
        """Initialize the ledger (create table if needed)."""
        if self._initialized:
            return True

        if not self.is_postgres_available:
            self._initialized = True
            return False  # Using memory fallback

        try:
            self._pool = await asyncpg.create_pool(
                self.database_url, min_size=1, max_size=5, command_timeout=30
            )
            async with self._pool.acquire() as conn:
                await conn.execute(self.CREATE_TABLE_SQL)
                # Get last hash for chain continuity
                row = await conn.fetchrow(
                    "SELECT entry_hash FROM vault999 ORDER BY id DESC LIMIT 1"
                )
                if row:
                    self._last_hash = row["entry_hash"]
            self._initialized = True
            return True
        except Exception as e:
            print(f"[SessionLedger] Postgres init failed: {e}, using memory fallback")
            self._initialized = True
            return False

    def _compute_hash(self, entry: VaultEntry) -> str:
        """Compute SHA256 hash for Merkle chain."""
        data = f"{entry.prev_hash}:{entry.session_id}:{entry.timestamp}:{entry.verdict_type}:{json.dumps(entry.payload, sort_keys=True)}"
        return hashlib.sha256(data.encode()).hexdigest()[:16]

    def _generate_entry_id(self) -> str:
        """Generate unique entry ID."""
        import secrets

        return secrets.token_hex(8)

    async def seal(
        self,
        session_id: str,
        verdict_type: str,
        payload: Dict[str, Any],
        query_summary: str = "",
        risk_level: str = "LOW",
        category: str = "general",
        environment: str = "prod",
        floors_checked: Optional[List[str]] = None,
        floors_failed: Optional[List[str]] = None,
    ) -> VaultEntry:
        """
        Seal a new entry into the VAULT999 ledger.

        Returns the sealed VaultEntry with computed hash.
        """
        await self.initialize()

        entry = VaultEntry(
            entry_id=self._generate_entry_id(),
            session_id=session_id,
            timestamp=datetime.now(timezone.utc).isoformat(),
            verdict_type=verdict_type,
            risk_level=risk_level,
            category=category,
            environment=environment,
            payload=payload,
            query_summary=query_summary,
            prev_hash=self._last_hash,
            entry_hash="",  # Computed below
            floors_checked=floors_checked or [],
            floors_failed=floors_failed or [],
        )

        entry.entry_hash = self._compute_hash(entry)
        self._last_hash = entry.entry_hash

        # Persist to Postgres or memory
        if self._pool:
            try:
                async with self._pool.acquire() as conn:
                    await conn.execute(
                        """
                        INSERT INTO vault999 (
                            entry_id, session_id, timestamp, verdict_type, risk_level,
                            category, environment, payload, query_summary,
                            prev_hash, entry_hash, schema_version, floors_checked, floors_failed
                        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
                        """,
                        entry.entry_id,
                        entry.session_id,
                        datetime.fromisoformat(entry.timestamp),
                        entry.verdict_type,
                        entry.risk_level,
                        entry.category,
                        entry.environment,
                        json.dumps(entry.payload),
                        entry.query_summary,
                        entry.prev_hash,
                        entry.entry_hash,
                        entry.schema_version,
                        entry.floors_checked,
                        entry.floors_failed,
                    )
            except Exception as e:
                print(f"[SessionLedger] Postgres write failed: {e}")
                self._memory_ledger.append(entry)
        else:
            self._memory_ledger.append(entry)

        return entry

    async def query(
        self,
        session_id: Optional[str] = None,
        verdict_type: Optional[str] = None,
        limit: int = 10,
        offset: int = 0,
    ) -> List[VaultEntry]:
        """Query the VAULT999 ledger."""
        await self.initialize()

        if self._pool:
            try:
                async with self._pool.acquire() as conn:
                    conditions = []
                    params = []
                    param_idx = 1

                    if session_id:
                        conditions.append(f"session_id = ${param_idx}")
                        params.append(session_id)
                        param_idx += 1

                    if verdict_type:
                        conditions.append(f"verdict_type = ${param_idx}")
                        params.append(verdict_type)
                        param_idx += 1

                    where_clause = " AND ".join(conditions) if conditions else "TRUE"

                    rows = await conn.fetch(
                        f"""
                        SELECT * FROM vault999 
                        WHERE {where_clause}
                        ORDER BY timestamp DESC
                        LIMIT ${param_idx} OFFSET ${param_idx + 1}
                        """,
                        *params,
                        limit,
                        offset,
                    )

                    return [
                        VaultEntry(
                            entry_id=row["entry_id"],
                            session_id=row["session_id"],
                            timestamp=row["timestamp"].isoformat() if row["timestamp"] else "",
                            verdict_type=row["verdict_type"],
                            risk_level=row["risk_level"],
                            category=row["category"] or "",
                            environment=row["environment"] or "",
                            payload=json.loads(row["payload"]) if row["payload"] else {},
                            query_summary=row["query_summary"] or "",
                            prev_hash=row["prev_hash"] or "",
                            entry_hash=row["entry_hash"],
                            schema_version=row["schema_version"] or "3.0",
                            floors_checked=list(row["floors_checked"] or []),
                            floors_failed=list(row["floors_failed"] or []),
                        )
                        for row in rows
                    ]
            except Exception as e:
                print(f"[SessionLedger] Postgres query failed: {e}")

        # Memory fallback query
        entries = self._memory_ledger
        if session_id:
            entries = [e for e in entries if e.session_id == session_id]
        if verdict_type:
            entries = [e for e in entries if e.verdict_type == verdict_type]
        return entries[offset : offset + limit]

    async def close(self):
        """Close the database connection pool."""
        if self._pool:
            await self._pool.close()
            self._pool = None


# Global ledger instance
_ledger: Optional[SessionLedger] = None


async def get_ledger() -> SessionLedger:
    """Get or create the global SessionLedger instance."""
    global _ledger
    if _ledger is None:
        _ledger = SessionLedger()
        await _ledger.initialize()
    return _ledger


async def seal_memory(
    session_id: str, verdict: str, payload: Dict[str, Any], **kwargs
) -> Dict[str, Any]:
    """
    Legacy compatibility function for APEX engine.
    Uses the new SessionLedger under the hood.
    """
    ledger = await get_ledger()
    entry = await ledger.seal(
        session_id=session_id, verdict_type=verdict, payload=payload, **kwargs
    )
    return {
        "verdict": entry.verdict_type,
        "seal_id": entry.entry_id,
        "entry_hash": entry.entry_hash,
    }


async def log_asi_decision(
    session_id: str,
    stage: str,
    query: str,
    asi_output: Dict[str, Any],
    verdict: str = "SEAL",
    floors_checked: Optional[List[str]] = None,
    floors_failed: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Log an ASI decision to VAULT for future fine-tuning (Ω Incident Logging).

    Args:
        session_id: Constitutional session token
        stage: "555" (empathy) or "666" (align)
        query: Original user query (will be anonymized)
        asi_output: ASI organ output (must include floor scores, stakeholder impact, metrics)
        verdict: Verdict from ASI stage (SEAL, SABAR, VOID)
        floors_checked: List of floors checked (optional)
        floors_failed: List of floors failed (optional)

    Returns:
        dict with seal_id and entry_hash
    """
    ledger = await get_ledger()

    # Anonymize query: store only hash and length
    query_hash = hashlib.sha256(query.encode()).hexdigest()[:16]
    query_len = len(query)

    # Prepare payload with ASI decision data
    payload = {
        "stage": stage,
        "query_hash": query_hash,
        "query_length": query_len,
        "asi_output": asi_output,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

    # Determine risk level based on verdict and floor failures
    risk_level = "LOW"
    if verdict == "VOID":
        risk_level = "HIGH"
    elif verdict == "SABAR":
        risk_level = "MEDIUM"

    # Determine category
    category = "asi_decision"
    if stage == "555":
        category = "asi_empathy"
    elif stage == "666":
        category = "asi_align"

    checked = floors_checked or []
    failed = floors_failed or []
    entry = await ledger.seal(
        session_id=session_id,
        verdict_type=verdict,
        payload=payload,
        query_summary=f"{stage} ASI decision - query hash {query_hash}",
        risk_level=risk_level,
        category=category,
        environment="prod",
        floors_checked=checked,
        floors_failed=failed,
    )

    return {
        "verdict": entry.verdict_type,
        "seal_id": entry.entry_id,
        "entry_hash": entry.entry_hash,
    }


async def inject_memory(session_id: str, context: Dict[str, Any]) -> bool:
    """
    Legacy compatibility function for memory injection.
    In VAULT999 v3, context is stored as part of the ledger entry payload.
    """
    # In the new architecture, memory injection is handled during seal
    # This function returns True for API compatibility
    return True
