"""
RAG Retrieval — generated by aiwg nlp new
Pattern: rag-pipeline
Dependencies: anthropic, openai (for embeddings), numpy

Install: pip install anthropic openai numpy
"""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any

import numpy as np
import openai

EMBEDDING_MODEL = "text-embedding-3-small"
CHUNK_SIZE = 512
CHUNK_OVERLAP = 64
RETRIEVAL_K = 5


# ---------------------------------------------------------------------------
# Chunking
# ---------------------------------------------------------------------------

def chunk_text(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
    """Split text into overlapping chunks by word count."""
    words = text.split()
    chunks = []
    step = size - overlap
    for i in range(0, len(words), step):
        chunk = " ".join(words[i : i + size])
        if chunk:
            chunks.append(chunk)
    return chunks


# ---------------------------------------------------------------------------
# Embeddings
# ---------------------------------------------------------------------------

def embed(texts: list[str]) -> list[list[float]]:
    """Embed a batch of texts using OpenAI embedding model."""
    client = openai.OpenAI()
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=texts)
    return [r.embedding for r in response.data]


def cosine_similarity(a: list[float], b: list[float]) -> float:
    va, vb = np.array(a), np.array(b)
    return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb) + 1e-10))


# ---------------------------------------------------------------------------
# Index (in-memory for dev; replace with vector store for production)
# ---------------------------------------------------------------------------

class InMemoryIndex:
    def __init__(self) -> None:
        self.chunks: list[str] = []
        self.embeddings: list[list[float]] = []

    def add_documents(self, documents: list[str]) -> None:
        chunks = []
        for doc in documents:
            chunks.extend(chunk_text(doc))
        self.chunks.extend(chunks)
        self.embeddings.extend(embed(chunks))

    def retrieve(self, query: str, k: int = RETRIEVAL_K) -> list[str]:
        if not self.chunks:
            return []
        query_embedding = embed([query])[0]
        scores = [cosine_similarity(query_embedding, e) for e in self.embeddings]
        top_k = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
        return [self.chunks[i] for i in top_k]


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

def build_context(retrieved_chunks: list[str]) -> str:
    """Format retrieved chunks as context for the generation prompt."""
    parts = [f"[Document {i+1}]\n{chunk}" for i, chunk in enumerate(retrieved_chunks)]
    return "\n\n---\n\n".join(parts)


# Instantiate index — load your documents here
_index = InMemoryIndex()


def load_documents(paths: list[Path]) -> None:
    """Load documents into the retrieval index."""
    docs = [p.read_text(encoding="utf-8") for p in paths]
    _index.add_documents(docs)


def retrieve(query: str, k: int = RETRIEVAL_K) -> str:
    """Retrieve relevant chunks for a query and format as context string."""
    chunks = _index.retrieve(query, k=k)
    return build_context(chunks)