"""
Supabase pgvector memory backend for cloud (Lambda) deployments.

Replaces S3MemoryBackend for all cloud scopes (both project and global).
Daily logs are still written to S3 via S3FileStorage — they are small,
sequential, and not vector data. Only the vector index lives in Supabase.

Required one-time SQL setup in Supabase (run in the SQL editor):

    CREATE EXTENSION IF NOT EXISTS vector;

    CREATE TABLE IF NOT EXISTS daita_org_memory (
        id               TEXT        PRIMARY KEY DEFAULT gen_random_uuid()::TEXT,
        org_id           INTEGER     NOT NULL REFERENCES organization(organization_id) ON DELETE CASCADE,
        project_name     TEXT,
        workspace        TEXT        NOT NULL,
        content          TEXT        NOT NULL,
        content_preview  TEXT        GENERATED ALWAYS AS (LEFT(content, 200)) STORED,
        memory_type      TEXT        NOT NULL DEFAULT 'long_term',
        category         TEXT,
        importance       NUMERIC(4,3) NOT NULL DEFAULT 0.500
                         CHECK (importance >= 0 AND importance <= 1),
        authoritative    BOOLEAN     NOT NULL DEFAULT FALSE,
        tags             TEXT[]      NOT NULL DEFAULT '{}',
        agent_id         TEXT,
        source           TEXT,
        pinned           BOOLEAN     NOT NULL DEFAULT FALSE,
        embedding        vector(1536),
        recall_count     INTEGER     NOT NULL DEFAULT 0,
        last_recalled_at TIMESTAMPTZ,
        content_hash     TEXT        GENERATED ALWAYS AS
                             (md5('direct:' || left(trim(content), 200))) STORED,
        created_at       TIMESTAMPTZ NOT NULL DEFAULT NOW(),
        updated_at       TIMESTAMPTZ NOT NULL DEFAULT NOW(),
        deleted_at       TIMESTAMPTZ
    );

    CREATE INDEX idx_org_memory_org_type
        ON daita_org_memory (org_id, memory_type)
        WHERE deleted_at IS NULL;

    CREATE INDEX idx_org_memory_workspace
        ON daita_org_memory (org_id, workspace)
        WHERE deleted_at IS NULL;

    CREATE INDEX idx_org_memory_project
        ON daita_org_memory (org_id, project_name, workspace)
        WHERE deleted_at IS NULL;

    CREATE INDEX idx_org_memory_authoritative
        ON daita_org_memory (org_id, authoritative)
        WHERE authoritative = TRUE AND deleted_at IS NULL;

    CREATE UNIQUE INDEX idx_org_memory_content_hash
        ON daita_org_memory (org_id, project_name, workspace, content_hash)
        WHERE deleted_at IS NULL;

    CREATE INDEX idx_org_memory_vector
        ON daita_org_memory
        USING hnsw (embedding vector_cosine_ops)
        WITH (m = 16, ef_construction = 64);

    CREATE OR REPLACE FUNCTION daita_recall_org_memory(
        p_org_id         INTEGER,
        p_embedding      vector(1536),
        p_limit          INTEGER  DEFAULT 5,
        p_threshold      FLOAT    DEFAULT 0.6,
        p_min_importance FLOAT    DEFAULT 0.0,
        p_workspace      TEXT     DEFAULT NULL,
        p_project_name   TEXT     DEFAULT NULL
    )
    RETURNS TABLE (
        id              TEXT,
        content         TEXT,
        content_preview TEXT,
        category        TEXT,
        importance      NUMERIC,
        authoritative   BOOLEAN,
        agent_id        TEXT,
        source          TEXT,
        pinned          BOOLEAN,
        recall_count    INTEGER,
        created_at      TIMESTAMPTZ,
        score           FLOAT
    )
    LANGUAGE plpgsql VOLATILE AS $$
    BEGIN
        SET LOCAL hnsw.ef_search = 100;
        RETURN QUERY
            SELECT
                m.id, m.content, m.content_preview, m.category,
                m.importance, m.authoritative, m.agent_id, m.source,
                m.pinned, m.recall_count, m.created_at,
                (1 - (m.embedding <=> p_embedding))::FLOAT AS score
            FROM daita_org_memory m
            WHERE m.org_id = p_org_id
              AND m.deleted_at IS NULL
              AND (p_project_name IS NULL OR m.project_name = p_project_name)
              AND (p_workspace IS NULL OR m.workspace = p_workspace)
              AND m.importance >= p_min_importance
              AND 1 - (m.embedding <=> p_embedding) >= p_threshold
            ORDER BY m.embedding <=> p_embedding
            LIMIT p_limit;
    END;
    $$;

    CREATE OR REPLACE FUNCTION daita_increment_recall_counts(
        p_ids         TEXT[],
        p_recalled_at TIMESTAMPTZ
    )
    RETURNS VOID
    LANGUAGE SQL AS $$
        UPDATE daita_org_memory
        SET recall_count     = recall_count + 1,
            last_recalled_at = p_recalled_at
        WHERE id = ANY(p_ids)
          AND deleted_at IS NULL;
    $$;

Required Lambda environment variables:
    SUPABASE_URL              - Supabase project URL (https://<ref>.supabase.co)
    SUPABASE_SERVICE_ROLE_KEY - Service role key (not the anon key)
    OPENAI_API_KEY            - For generating embeddings
"""

import os
import hashlib
import logging
import asyncio
from pathlib import Path
from datetime import datetime, timezone, date as _date
from typing import List, Dict, Any, Optional

from .metadata import MemoryMetadata

logger = logging.getLogger(__name__)


class SupabaseMemoryBackend:
    """
    Cloud memory backend using Supabase pgvector.

    Used for all Lambda deployments (scope="project" and scope="global").
    Replaces S3MemoryBackend entirely.

    Architecture:
    - Vectors: Supabase daita_org_memory table with HNSW index (O(log N) recall)
    - Daily logs: S3 via S3FileStorage (small sequential files, not vector data)
    - Writes: immediate ACID inserts, no batching/flush needed
    - Concurrency: safe under parallel Lambda invocations
    """

    def __init__(
        self,
        org_id: str,
        project_name: str,
        workspace: str,
        agent_id: str,
        scope: str = "project",
        embedding_provider: str = "openai",
        embedding_model: str = "text-embedding-3-small",
        supabase_url: Optional[str] = None,
        supabase_key: Optional[str] = None,
    ):
        self.org_id = int(org_id)
        self.project_name = project_name if scope != "global" else None
        self.workspace = workspace
        self.agent_id = agent_id
        self.scope = scope
        self.embedding_provider = embedding_provider
        self.embedding_model = embedding_model

        self._supabase_url = (
            supabase_url
            or os.getenv("SUPABASE_URL")
            or os.getenv("NEXT_PUBLIC_SUPABASE_URL")
        )
        self._supabase_key = (
            supabase_key
            or os.getenv("SUPABASE_SERVICE_ROLE_KEY")
        )

        if not self._supabase_url or not self._supabase_key:
            raise ValueError(
                "Supabase credentials required for cloud memory. "
                "Set SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY environment variables."
            )

        # Lazy async init — cannot call create_async_client in __init__
        self._client = None
        # In-session embedding cache — avoids duplicate API calls for the same text
        self._embed_cache: dict = {}

        # Embedding client
        if embedding_provider == "openai":
            from openai import AsyncOpenAI
            api_key = os.getenv("OPENAI_API_KEY")
            if not api_key:
                raise ValueError("OPENAI_API_KEY required for cloud memory embeddings")
            self._embedder = AsyncOpenAI(api_key=api_key)
        else:
            raise NotImplementedError(
                f"Embedding provider '{embedding_provider}' not supported. Use 'openai'."
            )

        # S3 storage for daily logs (logs are small, sequential, not vectors)
        import boto3
        s3 = boto3.client("s3")
        from .s3_storage import S3FileStorage
        bucket = f"daita-memory-{os.getenv('AWS_REGION', 'us-east-1')}"
        if scope == "global":
            log_prefix = f"orgs/{org_id}/global/workspaces/{workspace}/"
        else:
            log_prefix = f"orgs/{org_id}/projects/{project_name}/workspaces/{workspace}/"
        self.storage = S3FileStorage(s3, bucket, log_prefix, agent_id)

        # logs_dir as a Path so curator can do: backend.logs_dir / "2024-01-01.md"
        # S3FileStorage.read_file() handles relative paths by prepending the prefix
        self.logs_dir = Path("logs")

        # Metrics tracking
        self._pending_memory_count_increment = 0
        self._pending_memory_count_decrement = 0
        self._retrieval_count = 0

        logger.info(
            "SupabaseMemoryBackend initialized: org=%s, project=%s, workspace='%s', scope=%s",
            self.org_id, self.project_name, self.workspace, scope,
        )

    # ------------------------------------------------------------------
    # Supabase client — lazy async initialization
    # ------------------------------------------------------------------

    async def _get_client(self):
        if self._client is None:
            from supabase import create_async_client
            self._client = await create_async_client(
                self._supabase_url, self._supabase_key
            )
        return self._client

    # ------------------------------------------------------------------
    # Duck-type compatibility: search property
    # ------------------------------------------------------------------

    @property
    def search(self):
        """Duck-type compat — curator calls backend.search.store_chunk()."""
        return self

    # ------------------------------------------------------------------
    # Embedding
    # ------------------------------------------------------------------

    async def _embed(self, text: str) -> List[float]:
        if text in self._embed_cache:
            return self._embed_cache[text]
        response = await self._embedder.embeddings.create(
            input=text,
            model=self.embedding_model,
        )
        embedding = response.data[0].embedding
        self._embed_cache[text] = embedding
        return embedding

    # ------------------------------------------------------------------
    # Score adjustments (S-7): matches local backend behavior exactly
    # ------------------------------------------------------------------

    def _apply_score_adjustments(self, base_score: float, memory: dict) -> float:
        importance = float(memory.get("importance", 0.5))
        importance_boost = (importance - 0.5) * 0.2

        decay = 1.0
        if not memory.get("pinned", False):
            created_at_str = memory.get("created_at")
            if created_at_str:
                try:
                    created_at = datetime.fromisoformat(
                        str(created_at_str).replace("Z", "+00:00")
                    )
                    age_days = (datetime.now(timezone.utc) - created_at).days
                    decay = max(0.7, 1.0 - (age_days / 365) * 0.3)
                except Exception:
                    pass

        return min((base_score + importance_boost) * decay, 1.0)

    # ------------------------------------------------------------------
    # remember() — S-1: signature matches memory_plugin.py contract
    # ------------------------------------------------------------------

    async def remember(
        self,
        content: str,
        category: Optional[str] = None,
        metadata: Optional[MemoryMetadata] = None,
        tags: Optional[List[str]] = None,
        authoritative: bool = False,
    ) -> Dict[str, Any]:
        """
        Store a memory: writes to S3 daily log and inserts into Supabase pgvector.

        Signature matches the contract called by memory_plugin.py:
            await backend.remember(content, category=category, metadata=metadata)
        """
        if metadata is None:
            metadata = MemoryMetadata(
                content=content,
                importance=0.5,
                source="agent_inferred",
                category=category,
            )

        # S-3: Write to daily log so the curator has input for fact extraction
        await self.storage.append_to_daily_log(content, metadata.category)

        # Store in pgvector
        chunk_id = await self.store_chunk(
            content, metadata, authoritative=authoritative, tags=tags
        )

        return {
            "status": "success",
            "message": "Memory stored and indexed",
            "chunk_id": chunk_id,
            "indexed": True,
        }

    # ------------------------------------------------------------------
    # store_chunk() — S-6: curator write path, idempotent via content hash
    # ------------------------------------------------------------------

    async def store_chunk(
        self,
        content: str,
        metadata: MemoryMetadata,
        chunk_id: Optional[str] = None,
        authoritative: bool = False,
        tags: Optional[List[str]] = None,
    ) -> str:
        """
        Store a chunk in Supabase pgvector.

        Uses content-hash for idempotency — duplicate content is silently ignored
        at the database level via the UNIQUE index on content_hash.

        Returns:
            The database row ID of the stored (or existing) chunk
        """
        # P1-A: content-based hash (no timestamp), prevents intra-session duplicates
        content_hash = hashlib.md5(
            f"direct:{content.strip()[:200]}".encode()
        ).hexdigest()

        if chunk_id is None:
            chunk_id = content_hash

        final_importance = 1.0 if authoritative else metadata.importance
        embedding = await self._embed(content)
        client = await self._get_client()

        row = {
            "org_id": self.org_id,
            "project_name": self.project_name,
            "workspace": self.workspace,
            "content": content,
            "category": metadata.category,
            "importance": float(final_importance),
            "authoritative": authoritative,
            "tags": tags or [],
            "agent_id": self.agent_id,
            "source": metadata.source,
            "embedding": embedding,
            "pinned": metadata.pinned,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "updated_at": datetime.now(timezone.utc).isoformat(),
        }

        try:
            result = await (
                client.table("daita_org_memory")
                .insert(row)
                .execute()
            )
            inserted_id = result.data[0]["id"] if result.data else chunk_id
        except Exception as e:
            # Unique constraint violation = duplicate content, not an error
            if "duplicate" in str(e).lower() or "unique" in str(e).lower():
                logger.debug("Duplicate memory skipped: %s...", content[:60])
                return chunk_id
            raise

        self._pending_memory_count_increment += 1
        return inserted_id

    # ------------------------------------------------------------------
    # recall() — S-5: chunk_id in results, S-7: score adjustments
    # ------------------------------------------------------------------

    async def recall(
        self,
        query: str,
        limit: int = 5,
        score_threshold: float = 0.6,
        strategy: str = "semantic",  # Supabase always uses HNSW
        min_importance: Optional[float] = None,
        max_importance: Optional[float] = None,
        category: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Semantic recall via Supabase pgvector HNSW index.

        Always uses indexed vector search — O(log N), no full-table scan.
        Applies same importance boost and temporal decay as local backend.
        """
        embedding = await self._embed(query)
        client = await self._get_client()

        # Fetch extra rows — post-filter by adjusted score and category may reduce count
        loose_threshold = max(0.0, score_threshold - 0.15)
        fetch_limit = limit * 3 if category else limit * 2

        result = await client.rpc("daita_recall_org_memory", {
            "p_org_id": self.org_id,
            "p_embedding": embedding,
            "p_limit": fetch_limit,
            "p_threshold": loose_threshold,
            "p_min_importance": min_importance or 0.0,
            "p_workspace": self.workspace,
            "p_project_name": self.project_name,
        }).execute()

        memories = result.data or []

        # S-7: Apply score adjustments (importance boost + temporal decay)
        adjusted = []
        for m in memories:
            if category and m.get("category") != category:
                continue
            raw_score = float(m["score"])
            adj_score = self._apply_score_adjustments(raw_score, m)
            if adj_score >= score_threshold:
                adjusted.append({**m, "score": adj_score, "raw_semantic_score": raw_score})

        adjusted.sort(key=lambda x: x["score"], reverse=True)
        adjusted = adjusted[:limit]

        if max_importance is not None:
            adjusted = [m for m in adjusted if float(m["importance"]) <= max_importance]

        # Fire-and-forget recall count update
        if adjusted:
            ids = [m["id"] for m in adjusted]
            try:
                asyncio.create_task(self._increment_recall_counts(ids))
            except RuntimeError:
                pass  # No running event loop in edge cases
            self._retrieval_count += len(adjusted)

        # S-5: map id -> chunk_id, include raw_semantic_score (P3-2)
        return [
            {
                "content": m["content"],
                "chunk_id": m["id"],
                "score": float(m["score"]),
                "relevance_score": float(m["score"]),
                "raw_semantic_score": float(m["raw_semantic_score"]),
                "score_breakdown": {
                    "semantic": float(m["raw_semantic_score"]),
                },
                "metadata": {
                    "importance": float(m["importance"]),
                    "category": m.get("category"),
                    "pinned": m.get("pinned", False),
                    "access_count": m.get("recall_count", 0),
                    "created_at": str(m.get("created_at", "")),
                    "source": m.get("source"),
                    "authoritative": m.get("authoritative", False),
                },
            }
            for m in adjusted
        ]

    async def _increment_recall_counts(self, ids: List[str]):
        """Fire-and-forget: increment recall_count + last_recalled_at."""
        try:
            client = await self._get_client()
            await client.rpc("daita_increment_recall_counts", {
                "p_ids": ids,
                "p_recalled_at": datetime.now(timezone.utc).isoformat(),
            }).execute()
        except Exception as e:
            logger.debug("Failed to increment recall counts: %s", e)

    # ------------------------------------------------------------------
    # list_by_category() — direct dump, no embedding cost
    # ------------------------------------------------------------------

    async def list_by_category(
        self,
        category: str,
        min_importance: float = 0.0,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """Direct category dump — no embedding call, ordered by importance DESC."""
        client = await self._get_client()
        result = await (
            client.table("daita_org_memory")
            .select("id, content, category, importance, pinned, source, created_at")
            .eq("org_id", self.org_id)
            .eq("category", category)
            .is_("deleted_at", "null")
            .gte("importance", min_importance)
            .order("importance", desc=True)
            .limit(limit)
            .execute()
        )
        rows = result.data or []
        return [
            {
                "chunk_id": r["id"],
                "content": r["content"],
                "metadata": {
                    "importance": float(r.get("importance", 0.5)),
                    "category": r.get("category"),
                    "pinned": r.get("pinned", False),
                    "source": r.get("source"),
                    "created_at": str(r.get("created_at", "")),
                },
            }
            for r in rows
        ]

    # ------------------------------------------------------------------
    # track_access() — S-6: duck-type compat with backend.search.track_access
    # ------------------------------------------------------------------

    def track_access(self, chunk_ids: List[str]):
        """
        Increment access counts for recalled memories.
        Schedules a fire-and-forget coroutine — does not block the caller.
        """
        if not chunk_ids:
            return
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                asyncio.create_task(self._increment_recall_counts(chunk_ids))
        except Exception:
            pass

    # ------------------------------------------------------------------
    # update_memory() — S-6
    # ------------------------------------------------------------------

    async def update_memory(
        self,
        query: str,
        new_content: str,
        importance: float = 0.5,
    ) -> Dict[str, Any]:
        """
        Find memories matching query, soft-delete them, and store updated content.
        Uses soft delete to preserve the audit trail.
        """
        matches = await self.recall(query, limit=5, score_threshold=0.7)
        if not matches:
            return {"status": "not_found", "updated": 0, "message": "No matching memories found"}

        ids = [m["chunk_id"] for m in matches if m.get("chunk_id")]
        await self._soft_delete(ids)

        metadata = MemoryMetadata(
            content=new_content,
            importance=importance,
            source="agent_inferred",
        )
        await self.store_chunk(new_content, metadata)

        return {
            "status": "success",
            "updated": len(ids),
            "message": f"Replaced {len(ids)} memories with updated content",
        }

    async def _soft_delete(self, ids: List[str]):
        """Soft-delete rows by setting deleted_at."""
        if not ids:
            return
        client = await self._get_client()
        await (
            client.table("daita_org_memory")
            .update({"deleted_at": datetime.now(timezone.utc).isoformat()})
            .in_("id", ids)
            .eq("org_id", self.org_id)
            .execute()
        )
        self._pending_memory_count_decrement += len(ids)

    # ------------------------------------------------------------------
    # update_chunk_metadata() / delete_chunks() — P0-A
    # ------------------------------------------------------------------

    async def update_chunk_metadata(self, chunk_id: str, updates: dict):
        """Update metadata fields on a stored memory chunk."""
        client = await self._get_client()
        db_updates = {}
        if "importance" in updates:
            db_updates["importance"] = float(updates["importance"])
        if "pinned" in updates:
            db_updates["pinned"] = bool(updates["pinned"])
        if "category" in updates:
            db_updates["category"] = updates["category"]
        if db_updates:
            db_updates["updated_at"] = datetime.now(timezone.utc).isoformat()
            await (
                client.table("daita_org_memory")
                .update(db_updates)
                .eq("id", chunk_id)
                .eq("org_id", self.org_id)
                .execute()
            )

    async def delete_chunks(self, chunk_ids: List[str]):
        """Soft-delete chunks (preserves audit trail, allows rollback)."""
        await self._soft_delete(chunk_ids)

    # ------------------------------------------------------------------
    # regenerate_memory_md() — S-6
    # ------------------------------------------------------------------

    async def regenerate_memory_md(self, min_importance: float = 0.4) -> str:
        """
        Regenerate MEMORY.md by querying Supabase and writing the result to S3.
        Groups by category, sorted by importance descending. Curator-only write path.
        """
        client = await self._get_client()

        query = (
            client.table("daita_org_memory")
            .select("content, category, importance")
            .eq("org_id", self.org_id)
            .eq("workspace", self.workspace)
            .is_("deleted_at", "null")
            .gte("importance", min_importance)
            .order("importance", desc=True)
        )
        if self.project_name is not None:
            query = query.eq("project_name", self.project_name)
        else:
            query = query.is_("project_name", "null")

        result = await query.execute()
        rows = result.data or []

        by_category: Dict[str, List] = {}
        for row in rows:
            category = row.get("category") or "General"
            by_category.setdefault(category, []).append(
                (float(row["importance"]), row["content"].strip())
            )

        lines = ["# Long-Term Memory\n"]
        for category in sorted(by_category.keys()):
            entries = sorted(by_category[category], reverse=True)
            lines.append(f"\n## {category.title()}\n")
            for _, content in entries:
                lines.append(f"- {content}\n")

        memory_md = "".join(lines)

        # Write to S3 for accessibility via storage.read_file("MEMORY.md")
        import boto3 as _boto3
        s3 = _boto3.client("s3")
        s3.put_object(
            Bucket=self.storage.bucket,
            Key=f"{self.storage.prefix}MEMORY.md",
            Body=memory_md.encode("utf-8"),
            ContentType="text/markdown",
        )

        return f"{self.storage.prefix}MEMORY.md"

    # ------------------------------------------------------------------
    # read_memory_md() / read_today_log() — P0-B
    # ------------------------------------------------------------------

    async def read_memory_md(self) -> str:
        """Return MEMORY.md content (reads from S3)."""
        return await self.storage.read_file("MEMORY.md")

    async def read_today_log(self) -> str:
        """Return today's daily log content (reads from S3)."""
        today = _date.today().isoformat()
        return await self.storage.read_file(f"logs/{today}.md")

    async def read_memory(self, file_path: str) -> str:
        """Read an arbitrary memory file via S3 storage."""
        return await self.storage.read_file(file_path)

    # ------------------------------------------------------------------
    # flush() — no-op (Supabase writes are immediate)
    # ------------------------------------------------------------------

    async def flush(self):
        """No-op. Supabase writes are immediate — no flush required."""
        pass

    # ------------------------------------------------------------------
    # get_pending_metrics()
    # ------------------------------------------------------------------

    def get_pending_metrics(self) -> Dict[str, Any]:
        return {
            "memory_count_delta": (
                self._pending_memory_count_increment
                - self._pending_memory_count_decrement
            ),
            "memory_retrieval_count": self._retrieval_count,
        }
