embeddings/seed.py

import os
import re
import json
import time
import requests
from typing import List
from dotenv import load_dotenv

# Optional: Use supabase-py if available, else use requests
try:
    from supabase import create_client, Client
    try:
        from supabase import ClientOptions
    except ImportError:
        ClientOptions = None
    SUPABASE_AVAILABLE = True
except ImportError:
    SUPABASE_AVAILABLE = False
    ClientOptions = None

# Optional: Use BeautifulSoup for HTML parsing
try:
    from bs4 import BeautifulSoup
    SOUP_AVAILABLE = True
except ImportError:
    SOUP_AVAILABLE = False

# Optional: Use sentence-transformers for local embeddings
try:
    from sentence_transformers import SentenceTransformer
    SENTE_TRANSFORMERS_AVAILABLE = True
except Exception:
    SENTE_TRANSFORMERS_AVAILABLE = False
load_dotenv()

# Environment variables
SUPABASE_URL = os.getenv("SUPABASE_URL") or os.getenv("NEXT_PUBLIC_SUPABASE_URL")
SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
# provider keys
HF_TOKEN = os.getenv("HF_TOKEN")
SUPABASE_RAG_USER_ID = os.getenv("SUPABASE_RAG_USER_ID")

assert SUPABASE_URL, "SUPABASE_URL or NEXT_PUBLIC_SUPABASE_URL must be set in .env"
assert SUPABASE_SERVICE_ROLE_KEY, "SUPABASE_SERVICE_ROLE_KEY must be set in .env"
if not (HF_TOKEN or SENTE_TRANSFORMERS_AVAILABLE):
    raise AssertionError("Either HF_TOKEN or installed sentence-transformers is required for embeddings")

# Embedding model (default uses a Sentence-Transformers model unless overridden)
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
DEFAULT_HEADERS = {
    # Helps avoid anti-bot blocks when scraping public pages.
    "User-Agent": "Mozilla/5.0 (compatible; seed-script/1.0; +https://example.local)",
}


# Supabase client (type: Client | None)
REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "30"))
CONNECT_TIMEOUT = int(os.getenv("CONNECT_TIMEOUT", "10"))
EMBED_MAX_RETRIES = int(os.getenv("EMBED_MAX_RETRIES", "3"))
RETRY_BACKOFF_SECONDS = float(os.getenv("RETRY_BACKOFF_SECONDS", "1.5"))
if SUPABASE_AVAILABLE:
    if ClientOptions is not None:
        try:
            supabase = create_client(
                SUPABASE_URL,
                SUPABASE_SERVICE_ROLE_KEY,
                options=ClientOptions(postgrest_client_timeout=REQUEST_TIMEOUT),
            )
        except TypeError:
            supabase = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)
    else:
        supabase = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)
else:
    supabase = None

def split_text(text: str, chunk_size: int = 512, chunk_overlap: int = 100) -> List[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[start:end])
        start += chunk_size - chunk_overlap
    return chunks

def scrape_page(url: str) -> str:
    resp = requests.get(url, headers=DEFAULT_HEADERS, timeout=(CONNECT_TIMEOUT, REQUEST_TIMEOUT))
    resp.raise_for_status()
    html = resp.text
    if SOUP_AVAILABLE:
        soup = BeautifulSoup(html, "html.parser")
        return soup.get_text()
    else:
        return re.sub(r"<[^>]*>", "", html)

# PDF support disabled.

def embed_text(text: str) -> List[float]:
    """Create an embedding vector.

    Priority:
      1. Hugging Face Inference API (HF_TOKEN)
      2. Local sentence-transformers model (if installed)

    Retries transient errors according to configured limits.
    """
    last_error = None
    for attempt in range(1, EMBED_MAX_RETRIES + 1):
        try:
            # Provider selection (local model preferred)
            if SENTE_TRANSFORMERS_AVAILABLE:
                model = SentenceTransformer(EMBEDDING_MODEL)
                emb = model.encode(text, show_progress_bar=False, convert_to_numpy=True)
                embedding_list = list(emb.tolist() if hasattr(emb, "tolist") else emb)
                return [float(x) for x in embedding_list]
            elif HF_TOKEN:
                model_name = EMBEDDING_MODEL
                url = f"https://api-inference.huggingface.co/models/{model_name}"
                headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json", **DEFAULT_HEADERS}
                payload = {"inputs": text}
                resp = requests.post(url, headers=headers, json=payload, timeout=(CONNECT_TIMEOUT, REQUEST_TIMEOUT))
            else:
                raise RuntimeError("No embedding provider configured. Install sentence-transformers or set HF_TOKEN.")

            # Retry only for transient server/rate-limit failures.
            if resp.status_code == 429 or 500 <= resp.status_code < 600:
                raise requests.HTTPError(f"{resp.status_code}: {resp.text}", response=resp)

            resp.raise_for_status()
            j = resp.json()

            # Robust parsing for multiple providers' response formats
            embedding = None
            if isinstance(j, list):
                # HF Inference often returns a list of floats or list-of-lists
                if len(j) > 0 and isinstance(j[0], (float, int)):
                    embedding = j
                elif len(j) > 0 and isinstance(j[0], list):
                    try:
                        embedding = [sum(col) / len(j) for col in zip(*j)]
                    except Exception:
                        embedding = j[0]
            elif isinstance(j, dict):
                if "data" in j:
                    d = j["data"]
                    if isinstance(d, list) and len(d) > 0 and isinstance(d[0], dict):
                        embedding = d[0].get("embedding") or d[0].get("vector")
                    elif isinstance(d, list) and len(d) > 0 and isinstance(d[0], (list, float, int)):
                        embedding = d[0]
                    elif isinstance(d, dict):
                        embedding = d.get("embedding") or d.get("vector")
                elif "embedding" in j:
                    embedding = j.get("embedding")
                elif "result" in j and isinstance(j["result"], dict):
                    embedding = j["result"].get("embedding")
                elif "features" in j and isinstance(j["features"], list):
                    embedding = j["features"]
                elif "vector" in j:
                    embedding = j["vector"]

            if embedding is None:
                raise RuntimeError(f"Unexpected embedding response format: {j}")

            # Ensure the embedding is an iterable sequence and convert to floats
            if not isinstance(embedding, (list, tuple)):
                raise RuntimeError(f"Embedding value is not a sequence: {type(embedding)}")
            embedding_list = list(embedding)
            return [float(x) for x in embedding_list]

        except (requests.Timeout, requests.ConnectionError, requests.HTTPError) as exc:
            last_error = exc
            if attempt < EMBED_MAX_RETRIES:
                wait_s = RETRY_BACKOFF_SECONDS * attempt
                print(f"[seed] embed retry {attempt}/{EMBED_MAX_RETRIES} in {wait_s:.1f}s due to: {exc}")
                time.sleep(wait_s)
                continue
            raise RuntimeError(f"Embedding failed after {EMBED_MAX_RETRIES} attempts: {exc}") from exc

    # Defensive fallback; loop always returns or raises.
    raise RuntimeError(f"Embedding failed: {last_error}")

def insert_chunk_supabase(content: str, embedding: List[float], url: str):
    # Insert into public.rag_user_documents (preferred) or fallback REST path
    if supabase:
        data = {"content": content, "embedding": embedding, "url": url}
        if SUPABASE_RAG_USER_ID:
            data["user_id"] = SUPABASE_RAG_USER_ID
        res = supabase.table("rag_user_documents").insert(data).execute()
        # supabase-py returns an APIResponse with .data (list on success, dict on error)
        if hasattr(res, "data") and not isinstance(res.data, list):
            print(f"[seed] Error inserting chunk: {res.data}")
    else:
        # Fallback: direct REST API (if supabase-py not installed)
        headers = {
            "apikey": SUPABASE_SERVICE_ROLE_KEY,
            "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
            "Content-Type": "application/json"
        }
        data = {"content": content, "embedding": embedding, "url": url}
        if SUPABASE_RAG_USER_ID:
            data["user_id"] = SUPABASE_RAG_USER_ID
        url_api = f"{SUPABASE_URL}/rest/v1/rag_user_documents"
        resp = requests.post(url_api, headers=headers, data=json.dumps(data))
        if not resp.ok:
            print(f"[seed] Error inserting chunk: {resp.text}")

def clear_chunks_table():
    # Table clearing disabled — no-op to avoid accidental deletions.
    print("[seed] clear_chunks_table skipped (disabled).")

def insert_chunks(source: str, content: str):
    chunks = split_text(content)
    print(f"[seed] {source} split into {len(chunks)} chunks.")
    inserted = 0
    for idx, chunk in enumerate(chunks):
        try:
            vector = embed_text(chunk)
            insert_chunk_supabase(chunk, vector, source)
            inserted += 1
        except Exception as e:
            print(f"[seed] ERROR embedding/inserting chunk {idx + 1}/{len(chunks)} for {source}: {e}")
        if (idx + 1) % 10 == 0 or idx + 1 == len(chunks):
            print(f"[seed] {source} progress: {idx + 1}/{len(chunks)} chunks processed ({inserted} inserted).")
    print(f"[seed] Completed {source}: inserted {inserted}/{len(chunks)} chunks.")
    return inserted

def load_data(webpages: List[str]):
    print(f"[seed] Starting seed for {len(webpages)} webpages...")
    try:
        clear_chunks_table()
        print("[seed] Existing chunks cleared.")
    except Exception as e:
        print(f"[seed] ERROR: Failed to clear chunks table: {e}")
        print("[seed] Check that your Supabase project is active (free-tier projects pause after inactivity).")
        raise SystemExit(1)
    total_inserted = 0
    for idx, url in enumerate(webpages):
        print(f"[seed] ({idx + 1}/{len(webpages)}) Scraping: {url}")
        content = scrape_page(url)
        total_inserted += insert_chunks(url, content)
    print(f"[seed] Done. Total inserted chunks: {total_inserted}.")

if __name__ == "__main__":
    load_data([
        "https://jdentalspecialists.com/",
    ])