import os
from openai import OpenAI
from db import upsert_chunks
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    api_key=os.getenv("DEEPINFRA_API_KEY"),
    base_url="https://api.deepinfra.com/v1/openai",
)

CHUNK_SIZE    = 500
CHUNK_OVERLAP = 50
EMBEDDING_MODEL = "BAAI/bge-base-en-v1.5"
EMBEDDING_DIMS  = 768  # bge-base produces 768-dim vectors


def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunks.append(" ".join(words[start:end]))
        start += chunk_size - overlap
    return [c for c in chunks if c.strip()]


def embed_texts(texts: list[str]) -> list[list[float]]:
    all_embeddings = []
    batch_size = 100
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        response = client.embeddings.create(
            input=batch,
            model=EMBEDDING_MODEL,
        )
        all_embeddings.extend([item.embedding for item in response.data])
    return all_embeddings


def process_pages(filename: str, pages: list[dict],
                  nationality: str = "unknown", corpus: str = "unknown"):
    rows = []
    all_chunks = []

    for page in pages:
        chunks = chunk_text(page["text"])
        for idx, chunk in enumerate(chunks):
            all_chunks.append({
                "filename":    filename,
                "page_num":    page["page_num"],
                "chunk_index": idx,
                "text":        chunk,
                "nationality": nationality,
                "corpus":      corpus,
            })

    if not all_chunks:
        print(f"  No text extracted from {filename}, skipping.")
        return

    print(f"  Embedding {len(all_chunks)} chunks from {filename}...")
    texts = [c["text"] for c in all_chunks]
    embeddings = embed_texts(texts)

    for chunk, embedding in zip(all_chunks, embeddings):
        chunk["embedding"] = embedding
        rows.append(chunk)

    upsert_chunks(rows)
    print(f"  Stored {len(rows)} chunks for {filename}.")