import os, faiss
from pathlib import Path
from typing import List, Dict
from sentence_transformers import SentenceTransformer
import numpy as np

class CodebaseIndex:
    def __init__(self, project_path: str):
        self.project_path = Path(project_path)
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        self.index = None
        self.file_list = []
        self.chunks = []

    def build_index(self):
        exts = {'.py', '.js', '.ts', '.tsx', '.jsx', '.go', '.rs', '.java', '.cpp', '.c', '.h'}
        for filepath in self.project_path.rglob('*'):
            if filepath.suffix in exts and filepath.is_file():
                self.file_list.append(str(filepath.relative_to(self.project_path)))
                content = filepath.read_text()
                chunks = self._split_into_chunks(content)
                for chunk in chunks:
                    self.chunks.append((str(filepath.relative_to(self.project_path)), chunk))
        texts = [chunk[1] for chunk in self.chunks]
        embeddings = self.model.encode(texts, show_progress_bar=False)
        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dim)
        self.index.add(np.array(embeddings).astype('float32'))

    def search(self, query: str, top_k=5) -> List[Dict]:
        if not self.index:
            return []
        q_emb = self.model.encode([query]).astype('float32')
        D, I = self.index.search(q_emb, top_k)
        results = []
        for i, idx in enumerate(I[0]):
            if idx != -1:
                fpath, chunk = self.chunks[idx]
                results.append({"file": fpath, "snippet": chunk, "score": float(D[0][i])})
        return results

    def _split_into_chunks(self, text, max_chars=1000):
        paragraphs = text.split('\n\n')
        chunks = []
        current = ""
        for para in paragraphs:
            if len(current) + len(para) < max_chars:
                current += para + "\n\n"
            else:
                if current:
                    chunks.append(current.strip())
                current = para + "\n\n"
        if current:
            chunks.append(current.strip())
        return chunks