import numpy as np
from dataclasses import dataclass
from typing import List, Tuple


@dataclass
class Chunk:
    text: str
    source: str
    embedding: np.ndarray = None


def simple_bow_embed(text: str, vocab: dict) -> np.ndarray:
    """A minimal bag-of-words embedding stand-in (no external
    model dependency), sufficient to demonstrate the retrieval
    pipeline mechanics end to end without a real neural encoder."""
    vec = np.zeros(len(vocab))
    for word in text.lower().split():
        if word in vocab:
            vec[vocab[word]] += 1.0
    norm = np.linalg.norm(vec)
    return vec / norm if norm > 0 else vec


def build_vocab(texts: List[str]) -> dict:
    words = sorted(set(
        w for t in texts for w in t.lower().split()
    ))
    return {w: i for i, w in enumerate(words)}


def chunk_document(text: str, chunk_size: int = 30) -> List[str]:
    words = text.split()
    return [
        " ".join(words[i:i + chunk_size])
        for i in range(0, len(words), chunk_size)
    ]


class SimpleVectorIndex:
    def __init__(self, vocab):
        self.vocab = vocab
        self.chunks: List[Chunk] = []

    def add(self, text: str, source: str):
        emb = simple_bow_embed(text, self.vocab)
        self.chunks.append(Chunk(text, source, emb))

    def search(self, query: str, k: int = 3) -> List[Tuple[Chunk, float]]:
        q_emb = simple_bow_embed(query, self.vocab)
        scored = []
        for chunk in self.chunks:
            sim = float(np.dot(q_emb, chunk.embedding))
            scored.append((chunk, sim))
        scored.sort(key=lambda pair: -pair[1])
        return scored[:k]


def build_rag_prompt(query: str, retrieved: List[Tuple[Chunk, float]]) -> str:
    context = "\n\n".join(
        f"[Source: {c.source}]\n{c.text}" for c, _ in retrieved
    )
    return (
        "Answer the question using only the context below. "
        "If the context does not contain the answer, say so.\n\n"
        f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
    )
