backend/core/database.py

import os
from pathlib import Path
import re
from datetime import datetime
import hashlib
import uuid
from typing import Dict, Iterable, List, Optional, Any

from appwrite.client import Client
from appwrite.query import Query
from appwrite.services.databases import Databases
from dotenv import load_dotenv
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_qdrant import FastEmbedSparse, QdrantVectorStore, RetrievalMode
from qdrant_client import QdrantClient, AsyncQdrantClient, models
import asyncio

ROOT_ENV_PATH = Path(__file__).resolve().parents[2] / ".env"
if ROOT_ENV_PATH.exists():
    print(f"📡 Loading environment from: {ROOT_ENV_PATH}")
    load_dotenv(dotenv_path=ROOT_ENV_PATH, override=True)
else:
    print(f"⚠️ No .env file found at {ROOT_ENV_PATH}. Relying on system environment variables.")
    load_dotenv() # Fallback to standard search


def _env(name: str, default: Optional[str] = None) -> str:
    # Prioritise actual environment variables (set in HF Secrets)
    value = os.environ.get(name)
    if value and value.strip():
        return value
        
    # Fallback to default
    return default or ""


def get_qdrant_client() -> QdrantClient:
    url = _env("QDRANT_URL", "")
    api_key = _env("QDRANT_API_KEY", "")
    return QdrantClient(url=url, api_key=api_key, check_compatibility=False)

def get_async_qdrant_client() -> AsyncQdrantClient:
    url = _env("QDRANT_URL", "")
    api_key = _env("QDRANT_API_KEY", "")
    return AsyncQdrantClient(
        url=url,
        api_key=api_key,
        check_compatibility=False,
    )


class LocalSentenceTransformerEmbeddings(Embeddings):
    """LangChain-compatible local embeddings backed by sentence-transformers."""

    _model_cache: Dict[str, Any] = {}
    _pool_cache: Dict[str, Any] = {}

    def __init__(self, model_name_or_path: str) -> None:
        self.model_name_or_path = model_name_or_path

    def _get_model(self):
        cached = self._model_cache.get(self.model_name_or_path)
        if cached is not None:
            return cached

        from sentence_transformers import SentenceTransformer

        device = os.getenv("LOCAL_EMBED_DEVICE", "auto").strip().lower()
        model_kwargs = {}
        if device not in ("", "auto"):
            model_kwargs["device"] = device

        try:
            model = SentenceTransformer(self.model_name_or_path, **model_kwargs)
            if model_kwargs:
                print(f"⚡ Local embeddings using device: {device}")
        except Exception:
            # Fall back to auto device if the requested device backend is unavailable.
            model = SentenceTransformer(self.model_name_or_path)
            if model_kwargs:
                print(f"⚠ Requested device '{device}' unavailable. Falling back to auto device.")

        try:
            import torch

            num_threads = int(os.getenv("LOCAL_EMBED_NUM_THREADS", "0"))
            if num_threads > 0:
                torch.set_num_threads(num_threads)
        except Exception:
            pass

        self._model_cache[self.model_name_or_path] = model
        return model

    def _get_pool(self):
        cached = self._pool_cache.get(self.model_name_or_path)
        if cached is not None:
            return cached

        model = self._get_model()
        try:
            pool = model.start_multi_process_pool()
            self._pool_cache[self.model_name_or_path] = pool
            print("⚡ Local embedding multi-process pool enabled")
            return pool
        except Exception:
            return None

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        model = self._get_model()
        batch_size = int(os.getenv("LOCAL_EMBED_BATCH_SIZE", "256"))
        parallel = os.getenv("LOCAL_EMBED_PARALLEL", "true").strip().lower() in (
            "1",
            "true",
            "yes",
            "y",
        )
        min_parallel_docs = int(os.getenv("LOCAL_EMBED_PARALLEL_MIN_DOCS", "256"))

        vectors = None
        if parallel and len(texts) >= min_parallel_docs:
            pool = self._get_pool()
            if pool is not None:
                try:
                    vectors = model.encode_multi_process(
                        texts,
                        pool,
                        batch_size=batch_size,
                    )
                except Exception:
                    vectors = None

        if vectors is None:
            vectors = model.encode(
                texts,
                normalize_embeddings=True,
                batch_size=batch_size,
                show_progress_bar=False,
            )
        return vectors.tolist()

    def embed_query(self, text: str) -> List[float]:
        model = self._get_model()
        vector = model.encode(text, normalize_embeddings=True)
        return vector.tolist()


def _gemini_embeddings(model_candidates: List[str], api_key: Optional[str] = None) -> Embeddings:
    actual_key = (api_key or "").strip() or _env("GOOGLE_API_KEY", "").strip() or _env("GEMINI_API_KEY", "").strip()
    if not actual_key:
        err_msg = "GOOGLE_API_KEY is not set in environment or secrets. Please add it to your .env or Hugging Face Space Secrets."
        print(f"❌ {err_msg}")
        raise ValueError(err_msg)
        
    last_error = None
    for model_name in model_candidates:
        try:
            emb = GoogleGenerativeAIEmbeddings(
                model=model_name,
                google_api_key=actual_key,
            )
            # Minimal probe to verify key
            return emb
        except Exception as exc:
            last_error = exc
    raise RuntimeError("Unable to initialize Gemini embeddings") from last_error


def get_dense_embeddings(api_key: Optional[str] = None) -> Embeddings:
    provider = os.getenv("THEORY_EMBED_PROVIDER", "gemini").strip().lower()
    if provider in ("local", "sentence-transformers", "hf"):
        model_name = os.getenv(
            "THEORY_EMBEDDING_MODEL",
            "models/embeddings/bge-base-en-v1.5",
        )
        return LocalSentenceTransformerEmbeddings(model_name)

    return _gemini_embeddings(
        [
            os.getenv("THEORY_EMBEDDING_MODEL", "gemini-embedding-001"),
            "models/gemini-embedding-001",
        ],
        api_key=api_key
    )


def get_code_embeddings(api_key: Optional[str] = None) -> Embeddings:
    provider = os.getenv("CODE_EMBED_PROVIDER", "gemini").strip().lower()

    if provider in ("local", "sentence-transformers", "hf"):
        model_name = os.getenv(
            "CODE_EMBEDDING_MODEL",
            "models/embeddings/bge-base-en-v1.5",
        )
        return LocalSentenceTransformerEmbeddings(model_name)

    if provider == "voyage":
        try:
            from langchain_voyageai import VoyageAIEmbeddings

            model_name = os.getenv("CODE_EMBEDDING_MODEL", "voyage-code-3")
            emb = VoyageAIEmbeddings(model=model_name, voyage_api_key=api_key or os.getenv("VOYAGE_API_KEY"))
            emb.embed_query("module dff(input clk, input d, output reg q);")
            return emb
        except Exception:
            pass

    if provider == "openai":
        try:
            from langchain_openai import OpenAIEmbeddings

            model_name = os.getenv("CODE_EMBEDDING_MODEL", "text-embedding-3-large")
            emb = OpenAIEmbeddings(model=model_name, openai_api_key=api_key or os.getenv("OPENAI_API_KEY"))
            emb.embed_query("module dff(input clk, input d, output reg q);")
            return emb
        except Exception:
            pass

    return _gemini_embeddings(
        [
            os.getenv("CODE_EMBEDDING_MODEL", "text-embedding-004"),
            "models/text-embedding-004",
            "gemini-embedding-001",
        ],
        api_key=api_key
    )


def get_collection_name(kind: str) -> str:
    if kind == "theory":
        return os.getenv("QDRANT_COLLECTION_THEORY", "nandly_hardware_theory")
    if kind == "code":
        return os.getenv(
            "QDRANT_COLLECTION_CODE",
            os.getenv("QDRANT_COLLECTION", "nandly_hardware_rag"),
        )
    return os.getenv("QDRANT_COLLECTION", "nandly_hardware_rag")


async def clear_ingestion_collections() -> None:
    client = get_async_qdrant_client()
    collections_info = await client.get_collections()
    existing = {c.name for c in collections_info.collections}
    targets = [get_collection_name("code"), get_collection_name("theory")]

    for collection_name in targets:
        if collection_name in existing:
            await client.delete_collection(collection_name=collection_name)
            print(f"🧹 Cleared collection: {collection_name}")
        else:
            print(f"ℹ Collection not found (skip clear): {collection_name}")


async def ensure_hybrid_collection(
    client: AsyncQdrantClient,
    collection_name: str,
    embedding_dimension: int,
    dense_vector_name: str = "dense",
    sparse_vector_name: str = "bm25",
) -> None:
    collections_info = await client.get_collections()
    collections = {c.name for c in collections_info.collections}
    if collection_name in collections:
        return

    await client.create_collection(
        collection_name=collection_name,
        vectors_config={
            dense_vector_name: models.VectorParams(
                size=embedding_dimension,
                distance=models.Distance.COSINE,
            )
        },
        sparse_vectors_config={
            sparse_vector_name: models.SparseVectorParams(
                index=models.SparseIndexParams(on_disk=False)
            )
        },
    )


def get_vector_store(
    collection_name: Optional[str] = None,
    kind: str = "code",
) -> QdrantVectorStore:
    collection = collection_name or get_collection_name(kind)
    client = get_async_qdrant_client()
    embeddings = get_dense_embeddings(api_key=None) if kind == "theory" else get_code_embeddings(api_key=None)
    
    return QdrantVectorStore(
        client=client,
        collection_name=collection,
        embedding=embeddings,
        sparse_embedding=FastEmbedSparse(model_name="Qdrant/bm25"),
        retrieval_mode=RetrievalMode.HYBRID if kind == "theory" else RetrievalMode.DENSE,
        vector_name="dense",
        sparse_vector_name="bm25",
        async_mode=True
    )


async def upsert_documents(
    documents: Iterable[Document],
    batch_size: int = 64,
    kind: Optional[str] = None,
) -> int:
    """Upsert documents with automatic retry on rate limit errors (async)."""
    from google.api_core.exceptions import ResourceExhausted
    
    docs: List[Document] = list(documents)
    if not docs:
        return 0

    batch_size = int(os.getenv("UPSERT_BATCH_SIZE", str(batch_size)))

    inferred_kind = kind
    if inferred_kind is None:
        source_type = str(docs[0].metadata.get("source_type", "")).lower()
        inferred_kind = "theory" if "book" in source_type or source_type == "theory" else "code"

    store = get_vector_store(kind=inferred_kind)
    total = 0
    embed_provider = (
        os.getenv("THEORY_EMBED_PROVIDER", "gemini").strip().lower()
        if inferred_kind == "theory"
        else os.getenv("CODE_EMBED_PROVIDER", "gemini").strip().lower()
    )
    is_local_embeddings = embed_provider in ("local", "sentence-transformers", "hf")
    
    def _doc_id(doc: Document) -> str:
        meta = doc.metadata or {}
        stable_key = (
            str(meta.get("child_id") or "")
            or str(meta.get("record_id") or "")
            or str(meta.get("parent_id") or "")
            or f"{meta.get('source', '')}::{doc.page_content[:120]}"
        )
        digest = hashlib.sha1(stable_key.encode("utf-8", errors="ignore")).hexdigest()
        return str(uuid.uuid5(uuid.NAMESPACE_DNS, digest))

    for i in range(0, len(docs), batch_size):
        chunk = docs[i : i + batch_size]
        chunk_ids = [_doc_id(d) for d in chunk]
        max_retries = int(os.getenv("MAX_EMBED_RETRIES", "8"))
        retry_count = 0
        
        while retry_count < max_retries:
            try:
                await store.aadd_documents(chunk, ids=chunk_ids)
                total += len(chunk)
                print(f"✓ Indexed batch {i//batch_size + 1}: {len(chunk)} documents (total: {total}/{len(docs)})")
                
                if (not is_local_embeddings) and i + batch_size < len(docs):
                    await asyncio.sleep(1)
                break
                
            except ResourceExhausted as e:
                retry_count += 1
                error_msg = str(e)
                retry_match = re.search(r'retry in ([\d.]+)s', error_msg)
                wait_time = float(retry_match.group(1)) + 1 if retry_match else 60
                
                if retry_count < max_retries:
                    print(f"⚠ Rate limit hit. Waiting {wait_time:.1f}s before retry {retry_count}/{max_retries}...")
                    await asyncio.sleep(wait_time)
                else:
                    raise
                    
            except Exception as e:
                retry_count += 1
                error_msg = str(e)
                lowered = error_msg.lower()
                is_quota_error = any(x in lowered for x in ["429", "resource_exhausted", "quota", "rate limit", "retry in"])
                is_timeout_error = any(x in lowered for x in ["timed out", "timeout", "read operation timed out"])

                if is_quota_error and retry_count < max_retries:
                    retry_match = re.search(r"retry in ([\d.]+)s", error_msg, re.IGNORECASE)
                    wait_time = float(retry_match.group(1)) + 1 if retry_match else 60
                    print(f"⚠ Quota/rate-limit error. Waiting {wait_time:.1f}s before retry {retry_count}/{max_retries}...")
                    await asyncio.sleep(wait_time)
                    continue

                if is_timeout_error and retry_count < max_retries:
                    wait_time = int(os.getenv("TIMEOUT_RETRY_DELAY_SECONDS", "10"))
                    print(f"⚠ Timeout talking to Qdrant. Waiting {wait_time}s before retry {retry_count}/{max_retries}...")
                    await asyncio.sleep(wait_time)
                    continue

                print(f"✗ Error indexing batch: {e}")
                raise RuntimeError(f"Upsert failed after processing {total}/{len(docs)} documents: {e}") from e
    
    return total


# ======== Appwrite Chat History — Session-Document Model ========
# Each session = 1 row in chat_sessions, keyed by (userId, sessionId).
# Messages are serialised as a JSON blob inside the row.
# This avoids per-message rows and keeps all queries lightning fast.


def _get_tables_db():
    """Get Appwrite Databases client for chat operations."""
    client = Client()
    # Support both regional and universal endpoints
    endpoint = os.getenv("APPWRITE_ENDPOINT", "https://cloud.appwrite.io/v1")
    client.set_endpoint(endpoint)
    client.set_project(_env("APPWRITE_PROJECT_ID", "69afae6a000b5f5245c9"))
    client.set_key(_env("APPWRITE_API_KEY", ""))
    return Databases(client)


def _db_id() -> str:
    return _env("APPWRITE_DATABASE_ID", "69ce0fef002b79da9423")


_SESSIONS_TABLE = "chat_sessions"


import json as _json


async def create_chat_session(
    user_id: str,
    session_id: str,
    title: str = "New Chat",
) -> Dict:
    """Create a new chat session row in Appwrite."""
    try:
        db = _get_tables_db()
        now = datetime.utcnow().isoformat() + "Z"
        
        # Wrapped in to_thread for non-blocking sync SDK call
        row = await asyncio.to_thread(
            db.create_row,
            database_id=_db_id(),
            table_id=_SESSIONS_TABLE,
            row_id=session_id,
            data={
                "userId": user_id,
                "title": title,
                "messages": "[]",
                "isPinned": False,
                "lastUpdated": now,
                "createdAt": now,
            },
        )
        return row
    except Exception as e:
        print(f"Error creating chat session: {e}")
        raise


async def append_message_to_session(
    user_id: str,
    session_id: str,
    role: str,
    content: str,
    title: Optional[str] = None,
) -> Dict:
    """Append a message to a session (async)."""
    db = _get_tables_db()
    database_id = _db_id()

    try:
        row = await asyncio.to_thread(db.get_row, database_id, _SESSIONS_TABLE, session_id)
    except Exception:
        row = await create_chat_session(user_id, session_id, title or "New Chat")

    if row.get("userId") != user_id:
        raise PermissionError("Session does not belong to this user")

    existing_raw = row.get("messages") or "[]"
    try:
        messages: list = _json.loads(existing_raw)
    except (TypeError, _json.JSONDecodeError):
        messages = []

    messages.append({
        "role": role,
        "content": content,
        "timestamp": datetime.utcnow().isoformat() + "Z",
    })

    update_data = {
        "messages": _json.dumps(messages),
        "lastUpdated": datetime.utcnow().isoformat() + "Z",
    }
    if title:
        update_data["title"] = title[:500]

    return await asyncio.to_thread(
        db.update_row,
        database_id=database_id,
        table_id=_SESSIONS_TABLE,
        row_id=session_id,
        data=update_data,
    )


async def load_chat_session(user_id: str, session_id: str) -> Dict:
    """Load a chat session (async)."""
    db = _get_tables_db()
    row = await asyncio.to_thread(db.get_row, _db_id(), _SESSIONS_TABLE, session_id)

    if row.get("userId") != user_id:
        raise PermissionError("Session does not belong to this user")

    raw = row.get("messages") or "[]"
    try:
        messages = _json.loads(raw)
    except (TypeError, _json.JSONDecodeError):
        messages = []

    return {
        "id": row["$id"],
        "userId": row["userId"],
        "title": row.get("title", "New Chat"),
        "messages": messages,
        "isPinned": row.get("isPinned", False),
        "lastUpdated": row.get("lastUpdated"),
        "createdAt": row.get("createdAt"),
    }


async def list_user_sessions(user_id: str, limit: int = 50) -> List[Dict]:
    """List all chat sessions for a user (async)."""
    db = _get_tables_db()
    try:
        result = await asyncio.to_thread(
            db.list_rows,
            database_id=_db_id(),
            table_id=_SESSIONS_TABLE,
            queries=[
                Query.equal("userId", user_id),
                Query.order_desc("lastUpdated"),
                Query.limit(limit),
            ],
        )
        sessions = []
        for row in result.get("documents", result.get("rows", [])):
            sessions.append({
                "id": row["$id"],
                "title": row.get("title", "New Chat"),
                "isPinned": row.get("isPinned", False),
                "lastUpdated": row.get("lastUpdated"),
                "createdAt": row.get("createdAt"),
                "messageCount": len(_json.loads(row.get("messages") or "[]")),
            })
        return sessions
    except Exception as e:
        print(f"Error listing user sessions: {e}")
        return []


async def update_session_metadata(
    user_id: str,
    session_id: str,
    title: Optional[str] = None,
    is_pinned: Optional[bool] = None,
) -> Dict:
    db = _get_tables_db()
    database_id = _db_id()

    row = await asyncio.to_thread(db.get_row, database_id, _SESSIONS_TABLE, session_id)
    if row.get("userId") != user_id:
        raise PermissionError("Session does not belong to this user")

    data = {"lastUpdated": datetime.utcnow().isoformat() + "Z"}
    if title is not None:
        data["title"] = title[:500]
    if is_pinned is not None:
        data["isPinned"] = is_pinned

    return await asyncio.to_thread(
        db.update_row,
        database_id=database_id,
        table_id=_SESSIONS_TABLE,
        row_id=session_id,
        data=data,
    )


async def delete_chat_session(user_id: str, session_id: str) -> bool:
    """Hard-delete a chat session (async)."""
    db = _get_tables_db()
    database_id = _db_id()

    try:
        row = await asyncio.to_thread(db.get_row, database_id, _SESSIONS_TABLE, session_id)
    except Exception:
        return True

    if row.get("userId") != user_id:
        raise PermissionError("Session does not belong to this user")

    await asyncio.to_thread(
        db.delete_row,
        database_id=database_id,
        table_id=_SESSIONS_TABLE,
        row_id=session_id,
    )
    return True


async def update_full_session(
    user_id: str,
    session_id: str,
    title: Optional[str] = None,
    messages: Optional[list] = None,
    is_pinned: Optional[bool] = None,
) -> Dict:
    db = _get_tables_db()
    database_id = _db_id()

    try:
        row = await asyncio.to_thread(db.get_row, database_id, _SESSIONS_TABLE, session_id)
        if row.get("userId") != user_id:
            raise PermissionError("Session does not belong to this user")
    except PermissionError:
        raise
    except Exception:
        return await create_chat_session(user_id, session_id, title or "New Chat")

    data = {"lastUpdated": datetime.utcnow().isoformat() + "Z"}
    if title is not None:
        data["title"] = title[:500]
    if messages is not None:
        data["messages"] = _json.dumps(messages)
    if is_pinned is not None:
        data["isPinned"] = is_pinned

    return await asyncio.to_thread(
        db.update_row,
        database_id=database_id,
        table_id=_SESSIONS_TABLE,
        row_id=session_id,
        data=data,
    )


# ======== Legacy wrappers (backward compat for main.py) ========


async def save_chat_message(
    user_id: str,
    role: str,
    content: str,
    thread_id: str,
) -> Dict:
    """Legacy wrapper (async)."""
    return await append_message_to_session(
        user_id=user_id,
        session_id=thread_id,
        role=role,
        content=content,
    )


async def load_chat_history(
    user_id: str,
    thread_id: str,
    limit: int = 50,
) -> List[Dict]:
    """Legacy wrapper (async)."""
    try:
        session = await load_chat_session(user_id, thread_id)
        msgs = session.get("messages", [])
        return msgs[-limit:] if limit else msgs
    except Exception:
        return []


async def get_user_threads(user_id: str, limit: int = 20) -> List[Dict]:
    """Legacy wrapper (async)."""
    return await list_user_sessions(user_id, limit)


# ======== Appwrite User Stats ========


def get_appwrite_db_client() -> "Databases":
    """Get Appwrite Databases client for stats operations."""
    client = Client()
    # Explicit endpoint fallback
    endpoint = os.environ.get("APPWRITE_ENDPOINT") or "https://cloud.appwrite.io/v1"
    project = os.environ.get("APPWRITE_PROJECT_ID") or "69afae6a000b5f5245c9" # Your specific Project ID
    key = os.environ.get("APPWRITE_API_KEY") or ""
    
    client.set_endpoint(endpoint)
    client.set_project(project)
    if key:
        client.set_key(key)
        
    return Databases(client)


async def update_user_tokens(user_id: str, tokens_to_add: int) -> int:
    """Update total tokens used by a user in Appwrite (async)."""
    try:
        databases = get_appwrite_db_client()
        database_id = _env("APPWRITE_DATABASE_ID", "69ce0fef002b79da9423")
        collection_id = os.getenv("APPWRITE_USER_STATS_COLLECTION_ID", "user_stats")

        docs = await asyncio.to_thread(
            databases.list_rows,
            database_id=database_id,
            table_id=collection_id,
            queries=[Query.equal("userId", user_id)]
        )

        if docs["total"] > 0:
            doc = docs["documents"][0]
            new_total = (doc.get("totalTokens", 0) or 0) + tokens_to_add
            await asyncio.to_thread(
                databases.update_row,
                database_id=database_id,
                table_id=collection_id,
                row_id=doc["$id"],
                data={"totalTokens": new_total, "lastUsed": datetime.utcnow().isoformat()}
            )
            return new_total
        else:
            await asyncio.to_thread(
                databases.create_row,
                database_id=database_id,
                table_id=collection_id,
                row_id="unique()",
                data={
                    "userId": user_id,
                    "totalTokens": tokens_to_add,
                    "lastUsed": datetime.utcnow().isoformat()
                }
            )
            return tokens_to_add
    except Exception as e:
        print(f"Error updating user tokens: {e}")
        return 0


async def get_user_tokens(user_id: str) -> int:
    """Get total tokens used by a user (async)."""
    try:
        databases = get_appwrite_db_client()
        database_id = _env("APPWRITE_DATABASE_ID", "69ce0fef002b79da9423")
        collection_id = os.getenv("APPWRITE_USER_STATS_COLLECTION_ID", "user_stats")

        docs = await asyncio.to_thread(
            databases.list_rows,
            database_id=database_id,
            table_id=collection_id,
            queries=[Query.equal("userId", user_id)]
        )

        if docs["total"] > 0:
            return docs["documents"][0].get("totalTokens", 0)
        return 0
    except Exception as e:
        print(f"Error getting user tokens: {e}")
        return 0