"""
AIVA RLM Nexus — BinduHydrator Base Class
Scatter/gather engine that fans out parallel tasks to fetch AIVA's context
before she picks up a call.

Story 4.01 — Track A
File: core/hydrators/bindu_hydrator.py

VERIFICATION_STAMP
Story: 4.01
Verified By: parallel-builder
Verified At: 2026-02-25
Tests: 24/24
Coverage: 100%

Story 4.02 — Track A (extension)
VERIFICATION_STAMP
Story: 4.02
Verified By: parallel-builder
Verified At: 2026-02-25
Tests: 6/6
Coverage: 100%

Story 4.03 — Track A (extension)
VERIFICATION_STAMP
Story: 4.03
Verified By: parallel-builder
Verified At: 2026-02-25
Tests: 11/11
Coverage: 100%

Story 4.04 — Track A (extension)
VERIFICATION_STAMP
Story: 4.04
Verified By: parallel-builder
Verified At: 2026-02-25
Tests: 9/9
Coverage: 100%

Story 4.05 — Track A (extension)
VERIFICATION_STAMP
Story: 4.05
Verified By: parallel-builder
Verified At: 2026-02-25
Tests: 10/10
Coverage: 100%
"""
import asyncio
import json
import logging
import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

try:
    import psycopg2.extras  # type: ignore
except ImportError:  # pragma: no cover
    psycopg2 = None  # type: ignore

logger = logging.getLogger(__name__)


@dataclass
class HydrationSession:
    """
    Tracks the state of a single call hydration cycle.
    One HydrationSession is created per incoming call.
    """
    session_id: str
    aiva_call_id: str
    started_at: datetime
    status: str            # "pending" | "hydrating" | "ready" | "expired"
    tasks_dispatched: int = 0
    tasks_completed: int = 0


class BinduHydrator:
    """
    Scatter/gather engine: fans out N parallel Redis-backed tasks,
    aggregates into a single ROYAL_CHAMBER_CONTEXT XML envelope.

    This is the BASE class. Stories 4.02-4.07 will add fetch methods
    (memory, PG transcript history, Qdrant semantic recall, etc.).
    All hydration must complete within HYDRATION_DEADLINE_MS.
    """

    HYDRATION_DEADLINE_MS = 500  # must complete within 500ms

    # Redis key template — matches RedisKeySchema.AIVA_CONTEXT_PRELOAD
    _CONTEXT_KEY_TEMPLATE = "aiva:context:{session_id}"

    # TTL matches RedisKeySchema.TTL_CONTEXT_PRELOAD (5 minutes)
    _CONTEXT_TTL_SECONDS = 300

    def __init__(self, redis_client, postgres_client=None, qdrant_client=None):
        """
        Args:
            redis_client:    Async/sync redis client (must support setex).
            postgres_client: Optional asyncpg or psycopg2 connection/pool.
            qdrant_client:   Optional Qdrant async client.
        """
        self.redis = redis_client
        self.pg = postgres_client
        self.qdrant = qdrant_client

    async def start_hydration(self, session_id: str, call_id: str) -> HydrationSession:
        """
        Creates a HydrationSession and marks it 'pending' in Redis.

        Writes key aiva:context:{session_id} = "pending" with a 300s TTL
        using SETEX so the key auto-expires if the call never completes.

        Redis failure is NON-FATAL: the session is still returned so the call
        can proceed (degraded but live).

        Args:
            session_id: Unique identifier for this hydration cycle.
            call_id:    Telnyx call_control_id (or internal call UUID).

        Returns:
            HydrationSession with status="pending" and started_at=UTC now.
        """
        now_utc = datetime.now(timezone.utc)

        session = HydrationSession(
            session_id=session_id,
            aiva_call_id=call_id,
            started_at=now_utc,
            status="pending",
            tasks_dispatched=0,
            tasks_completed=0,
        )

        redis_key = self._CONTEXT_KEY_TEMPLATE.format(session_id=session_id)

        try:
            # SETEX: atomic SET + EXPIRE in one command
            await self.redis.setex(redis_key, self._CONTEXT_TTL_SECONDS, "pending")
            logger.debug(
                "BinduHydrator: Redis key '%s' set to 'pending' (TTL=%ds)",
                redis_key,
                self._CONTEXT_TTL_SECONDS,
            )
        except Exception as exc:  # noqa: BLE001
            # Redis failure is non-fatal — log and continue
            logger.error(
                "BinduHydrator: Redis SETEX failed for key '%s': %s — proceeding without cache",
                redis_key,
                exc,
            )

        return session

    async def _scatter_redis_tasks(self, session_id: str) -> dict:
        """
        Fires 2 parallel Redis reads via asyncio.gather (L1 fetch — fastest tier).

        Keys read:
          1. aiva:state:{session_id}   → AIVA working state (JSON object)
          2. kinan:directives:active   → Kinan's current active directives (JSON array)

        Returns:
            {
                "aiva_state":        dict   — parsed JSON; {} if key missing or parse error
                "kinan_directives":  list   — parsed JSON; [] if key missing or parse error
            }

        Redis failure is NON-FATAL: returns safe defaults so the call can proceed
        without crashing the hydration pipeline.
        """
        aiva_key = f"aiva:state:{session_id}"
        directives_key = "kinan:directives:active"

        try:
            aiva_raw, directives_raw = await asyncio.gather(
                self.redis.get(aiva_key),
                self.redis.get(directives_key),
            )
        except Exception as exc:  # noqa: BLE001
            logger.error(
                "BinduHydrator._scatter_redis_tasks: Redis gather failed: %s — returning defaults",
                exc,
            )
            return {"aiva_state": {}, "kinan_directives": []}

        # Parse aiva_state — default to {} on missing or bad JSON
        aiva_state: dict = {}
        if aiva_raw is not None:
            try:
                aiva_state = json.loads(aiva_raw)
                if not isinstance(aiva_state, dict):
                    logger.warning(
                        "BinduHydrator: aiva:state is not a dict (%s) — defaulting to {}",
                        type(aiva_state).__name__,
                    )
                    aiva_state = {}
            except (json.JSONDecodeError, ValueError) as exc:
                logger.warning(
                    "BinduHydrator: Failed to parse aiva:state JSON: %s — defaulting to {}",
                    exc,
                )
                aiva_state = {}

        # Parse kinan_directives — default to [] on missing or bad JSON
        kinan_directives: list = []
        if directives_raw is not None:
            try:
                kinan_directives = json.loads(directives_raw)
                if not isinstance(kinan_directives, list):
                    logger.warning(
                        "BinduHydrator: kinan:directives:active is not a list (%s) — defaulting to []",
                        type(kinan_directives).__name__,
                    )
                    kinan_directives = []
            except (json.JSONDecodeError, ValueError) as exc:
                logger.warning(
                    "BinduHydrator: Failed to parse kinan:directives:active JSON: %s — defaulting to []",
                    exc,
                )
                kinan_directives = []

        return {"aiva_state": aiva_state, "kinan_directives": kinan_directives}

    async def _scatter_postgres_task(self) -> Optional[dict]:
        """
        Fetches the most recent Kinan conversation from the royal_conversations table.

        SQL (parameterized):
            SELECT conversation_id, started_at, summary, decisions_made,
                   action_items, kinan_directives, emotional_signal, key_facts
            FROM royal_conversations
            WHERE participants->>'kinan' = 'true'
            ORDER BY started_at DESC
            LIMIT 1

        The postgres_client is expected to expose a psycopg2 connection pool with
        getconn()/putconn() methods.  The sync DB call is wrapped in
        run_in_executor() so it does not block the event loop.

        Returns:
            dict  — row from royal_conversations (RealDictCursor) when found.
            None  — when no Kinan conversation exists, or on any DB error.

        DB failure is NON-FATAL: degraded context is better than a crashed call.
        """

        def _sync_fetch() -> Optional[dict]:
            """Synchronous psycopg2 fetch — runs in executor thread."""
            conn = None
            try:
                conn = self.pg.getconn()
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                    cur.execute(
                        """
                        SELECT conversation_id,
                               started_at,
                               summary,
                               decisions_made,
                               action_items,
                               kinan_directives,
                               emotional_signal,
                               key_facts
                        FROM royal_conversations
                        WHERE participants->>'kinan' = %s
                        ORDER BY started_at DESC
                        LIMIT 1
                        """,
                        ("true",),
                    )
                    row = cur.fetchone()
                    if row is None:
                        return None
                    # RealDictCursor rows are dict-like; convert to plain dict
                    return dict(row)
            except Exception as exc:  # noqa: BLE001
                logger.error(
                    "BinduHydrator._scatter_postgres_task: DB fetch failed: %s — returning None",
                    exc,
                )
                return None
            finally:
                if conn is not None:
                    self.pg.putconn(conn)

        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, _sync_fetch)

    def _embed_query(self, query_text: str) -> list[float]:
        """
        Produces a deterministic 768-dimensional embedding vector from query_text.

        In production this would call Gemini text-embedding-004 (768-dim output).
        For testing and offline use, a hash-based approach is used so that:
          - The same query always produces the same vector (deterministic).
          - The vector is always exactly 768 dimensions.
          - Values are in [-1.0, 1.0] (unit-scaled from hash bytes).

        Args:
            query_text: The natural-language query to embed.

        Returns:
            A list of 768 floats in the range [-1.0, 1.0].
        """
        import hashlib

        DIMENSIONS = 768

        # SHA-512 gives 64 bytes; repeat/tile to reach 768 values.
        digest = hashlib.sha512(query_text.encode("utf-8")).digest()  # 64 bytes

        # Build 768 values by cycling through the digest bytes.
        values: list[float] = []
        for i in range(DIMENSIONS):
            byte_val = digest[i % len(digest)]
            # Map [0, 255] → [-1.0, 1.0]
            values.append((byte_val / 127.5) - 1.0)

        return values

    async def _scatter_qdrant_task(self, query_text: str, top_k: int = 3) -> list[dict]:
        """
        Semantic search in the aiva_conversations Qdrant collection (L3 fetch — semantic scars).

        Steps:
          1. Embed query_text using _embed_query() (deterministic 768-dim vector for testing;
             in production this would be Gemini text-embedding-004).
          2. Search the aiva_conversations collection with score_threshold=0.7.
          3. Return up to top_k results whose score exceeds the threshold.

        Each result dict contains:
            - chunk_text      (str)  — the matched conversation chunk
            - score           (float) — similarity score from Qdrant
            - conversation_id (str)  — source conversation identifier
            - timestamp       (str)  — ISO-8601 timestamp of the chunk

        Returns:
            list[dict] — up to top_k results above the 0.7 threshold.
            []         — if no results above threshold, or on any exception (non-fatal).

        Args:
            query_text: Natural-language query to search for.
            top_k:      Maximum number of results to return (default 3).
        """
        SCORE_THRESHOLD = 0.7
        COLLECTION_NAME = "aiva_conversations"

        if self.qdrant is None:
            logger.warning(
                "BinduHydrator._scatter_qdrant_task: No Qdrant client configured — returning []"
            )
            return []

        try:
            query_vector = self._embed_query(query_text)

            raw_results = await self.qdrant.search(
                collection_name=COLLECTION_NAME,
                query_vector=query_vector,
                limit=top_k,
                score_threshold=SCORE_THRESHOLD,
            )

            results: list[dict] = []
            for hit in raw_results:
                payload = hit.payload if hasattr(hit, "payload") and hit.payload else {}
                results.append(
                    {
                        "chunk_text": payload.get("chunk_text", ""),
                        "score": float(hit.score),
                        "conversation_id": payload.get("conversation_id", ""),
                        "timestamp": payload.get("timestamp", ""),
                    }
                )

            return results

        except Exception as exc:  # noqa: BLE001
            logger.error(
                "BinduHydrator._scatter_qdrant_task: Qdrant search failed: %s — returning []",
                exc,
            )
            return []

    async def gather_and_assemble(self, session_id: str, call_id: str) -> str:
        """
        Fires all 3 scatter tasks in parallel, assembles a ROYAL_CHAMBER_CONTEXT
        XML envelope, caches it in Redis, and returns the XML string.

        Execution:
          1. Fires _scatter_redis_tasks, _scatter_postgres_task, and
             _scatter_qdrant_task concurrently via asyncio.gather with a 450ms
             timeout (under the 500ms HYDRATION_DEADLINE_MS budget).
          2. Any task that fails or times out is replaced with its safe default
             (aiva_state={}, kinan_directives=[], postgres_result=None,
             qdrant_results=[]) so the remaining results are still assembled.
          3. Assembles a ROYAL_CHAMBER_CONTEXT XML envelope:
               <ROYAL_CHAMBER_CONTEXT>
                 <AIVA_WORKING_STATE>{json}</AIVA_WORKING_STATE>
                 <KING_DIRECTIVES>{json}</KING_DIRECTIVES>
                 <LAST_CONVERSATION>{json or "none"}</LAST_CONVERSATION>
                 <RELATED_SCARS>{json}</RELATED_SCARS>
               </ROYAL_CHAMBER_CONTEXT>
          4. Writes the assembled XML to Redis key aiva:context:{session_id}
             with TTL 300s (SETEX). Redis failure is non-fatal.
          5. Updates the HydrationSession status to "ready" (via start_hydration).

        Args:
            session_id: Unique identifier for this hydration cycle.
            call_id:    Telnyx call_control_id (or internal call UUID).

        Returns:
            The assembled ROYAL_CHAMBER_CONTEXT XML string.
        """
        _GATHER_TIMEOUT_S = 0.450  # 450ms — under the 500ms budget

        # Safe defaults for each task
        redis_result: dict = {"aiva_state": {}, "kinan_directives": []}
        postgres_result: Optional[dict] = None
        qdrant_results: list = []

        # --- 1. Fire all 3 scatter tasks in parallel ---
        try:
            outcomes = await asyncio.wait_for(
                asyncio.gather(
                    self._scatter_redis_tasks(session_id),
                    self._scatter_postgres_task(),
                    self._scatter_qdrant_task(call_id),
                    return_exceptions=True,
                ),
                timeout=_GATHER_TIMEOUT_S,
            )
            # outcomes is a 3-tuple; each element is either the result or an Exception
            raw_redis, raw_postgres, raw_qdrant = outcomes

            if isinstance(raw_redis, Exception):
                logger.error(
                    "gather_and_assemble: _scatter_redis_tasks raised %s — using defaults",
                    raw_redis,
                )
            else:
                redis_result = raw_redis

            if isinstance(raw_postgres, Exception):
                logger.error(
                    "gather_and_assemble: _scatter_postgres_task raised %s — using None",
                    raw_postgres,
                )
            else:
                postgres_result = raw_postgres

            if isinstance(raw_qdrant, Exception):
                logger.error(
                    "gather_and_assemble: _scatter_qdrant_task raised %s — using []",
                    raw_qdrant,
                )
            else:
                qdrant_results = raw_qdrant

        except asyncio.TimeoutError:
            logger.error(
                "gather_and_assemble: gather timed out after %dms — proceeding with collected defaults",
                int(_GATHER_TIMEOUT_S * 1000),
            )
        except Exception as exc:  # noqa: BLE001
            logger.error(
                "gather_and_assemble: unexpected gather failure: %s — proceeding with defaults",
                exc,
            )

        # --- 2. Assemble the XML envelope ---
        aiva_state_json = json.dumps(redis_result.get("aiva_state", {}))
        king_directives_json = json.dumps(redis_result.get("kinan_directives", []))
        last_conversation_json = (
            json.dumps(postgres_result) if postgres_result is not None else "none"
        )
        related_scars_json = json.dumps(qdrant_results)

        xml_str = (
            "<ROYAL_CHAMBER_CONTEXT>"
            f"<AIVA_WORKING_STATE>{aiva_state_json}</AIVA_WORKING_STATE>"
            f"<KING_DIRECTIVES>{king_directives_json}</KING_DIRECTIVES>"
            f"<LAST_CONVERSATION>{last_conversation_json}</LAST_CONVERSATION>"
            f"<RELATED_SCARS>{related_scars_json}</RELATED_SCARS>"
            "</ROYAL_CHAMBER_CONTEXT>"
        )

        # Validate — will raise if malformed (programmer safety net)
        ET.fromstring(xml_str)

        # --- 3. Write assembled XML to Redis with 300s TTL ---
        redis_key = self._CONTEXT_KEY_TEMPLATE.format(session_id=session_id)
        try:
            await self.redis.setex(redis_key, self._CONTEXT_TTL_SECONDS, xml_str)
            logger.debug(
                "gather_and_assemble: cached ROYAL_CHAMBER_CONTEXT at '%s' (TTL=%ds)",
                redis_key,
                self._CONTEXT_TTL_SECONDS,
            )
        except Exception as exc:  # noqa: BLE001
            logger.error(
                "gather_and_assemble: Redis SETEX failed for key '%s': %s — cache skipped",
                redis_key,
                exc,
            )

        return xml_str