"""
Core RapidRAG class - the main interface for local RAG.
"""

import os
import hashlib
from pathlib import Path
from typing import Optional, List, Dict, Any, Union
from datetime import datetime

import chromadb
from chromadb.config import Settings

# Optional imports
try:
    import httpx
    HTTPX_AVAILABLE = True
except ImportError:
    HTTPX_AVAILABLE = False


class RapidRAG:
    """
    Fast local RAG - search your documents with AI, no cloud needed.

    Example:
        rag = RapidRAG("my_project")
        rag.add("doc1", "Some text content")
        results = rag.search("query")
        answer = rag.query("What is...?", model="qwen2.5:7b")
    """

    def __init__(
        self,
        collection_name: str = "default",
        persist_dir: Optional[str] = None,
        embedding_model: str = "all-MiniLM-L6-v2",
        ollama_url: str = "http://localhost:11434",
    ):
        """
        Initialize RapidRAG.

        Args:
            collection_name: Name for the document collection
            persist_dir: Directory to persist the database (default: ./rapid_rag_data/)
            embedding_model: Sentence-transformers model for embeddings
            ollama_url: URL for Ollama API (for LLM queries)
        """
        self.collection_name = collection_name
        self.persist_dir = persist_dir or f"./rapid_rag_data/{collection_name}"
        self.embedding_model = embedding_model
        self.ollama_url = ollama_url

        # Create persist directory
        os.makedirs(self.persist_dir, exist_ok=True)

        # Initialize ChromaDB with persistence
        self.client = chromadb.PersistentClient(
            path=self.persist_dir,
            settings=Settings(anonymized_telemetry=False)
        )

        # Get or create collection
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}
        )

        # Lazy load embedding model
        self._embedder = None

    @property
    def embedder(self):
        """Lazy load sentence-transformers model."""
        if self._embedder is None:
            from sentence_transformers import SentenceTransformer
            self._embedder = SentenceTransformer(self.embedding_model)
        return self._embedder

    def _generate_id(self, content: str) -> str:
        """Generate a unique ID for content."""
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def add(
        self,
        doc_id: str,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Add a document to the collection.

        Args:
            doc_id: Unique document identifier
            content: Text content to index
            metadata: Optional metadata dict

        Returns:
            Document ID
        """
        metadata = metadata or {}
        metadata["added_at"] = datetime.now().isoformat()
        metadata["content_hash"] = self._generate_id(content)

        # Generate embedding
        embedding = self.embedder.encode(content).tolist()

        # Add to collection
        self.collection.add(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[content],
            metadatas=[metadata]
        )

        return doc_id

    def add_texts(
        self,
        texts: List[str],
        ids: Optional[List[str]] = None,
        metadatas: Optional[List[Dict]] = None
    ) -> List[str]:
        """
        Add multiple documents at once (faster).

        Args:
            texts: List of text contents
            ids: Optional list of IDs (auto-generated if not provided)
            metadatas: Optional list of metadata dicts

        Returns:
            List of document IDs
        """
        if ids is None:
            ids = [self._generate_id(t) for t in texts]

        if metadatas is None:
            metadatas = [{} for _ in texts]

        # Add timestamps
        now = datetime.now().isoformat()
        for meta in metadatas:
            meta["added_at"] = now

        # Generate embeddings (batch)
        embeddings = self.embedder.encode(texts).tolist()

        # Add to collection
        self.collection.add(
            ids=ids,
            embeddings=embeddings,
            documents=texts,
            metadatas=metadatas
        )

        return ids

    def add_file(
        self,
        file_path: Union[str, Path],
        chunk_size: int = 1000,
        chunk_overlap: int = 200
    ) -> List[str]:
        """
        Add a file to the collection.

        Supports: .txt, .md, .pdf (with pdf extra)

        Args:
            file_path: Path to file
            chunk_size: Characters per chunk
            chunk_overlap: Overlap between chunks

        Returns:
            List of chunk IDs
        """
        file_path = Path(file_path)

        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Read content based on file type
        suffix = file_path.suffix.lower()

        if suffix in [".txt", ".md"]:
            content = file_path.read_text(encoding="utf-8", errors="ignore")
        elif suffix == ".pdf":
            content = self._read_pdf(file_path)
        else:
            # Try reading as text
            content = file_path.read_text(encoding="utf-8", errors="ignore")

        # Chunk the content
        chunks = self._chunk_text(content, chunk_size, chunk_overlap)

        # Create IDs and metadata
        ids = [f"{file_path.stem}_{i}" for i in range(len(chunks))]
        metadatas = [
            {
                "source": str(file_path),
                "chunk_index": i,
                "total_chunks": len(chunks)
            }
            for i in range(len(chunks))
        ]

        return self.add_texts(chunks, ids, metadatas)

    def add_directory(
        self,
        dir_path: Union[str, Path],
        extensions: Optional[List[str]] = None,
        recursive: bool = True,
        **kwargs
    ) -> List[str]:
        """
        Add all files in a directory.

        Args:
            dir_path: Path to directory
            extensions: File extensions to include (default: [".txt", ".md", ".pdf"])
            recursive: Search subdirectories
            **kwargs: Passed to add_file()

        Returns:
            List of all chunk IDs
        """
        dir_path = Path(dir_path)
        extensions = extensions or [".txt", ".md", ".pdf"]

        all_ids = []
        pattern = "**/*" if recursive else "*"

        for ext in extensions:
            for file_path in dir_path.glob(f"{pattern}{ext}"):
                if file_path.is_file():
                    try:
                        ids = self.add_file(file_path, **kwargs)
                        all_ids.extend(ids)
                        print(f"✓ {file_path.name}: {len(ids)} chunks")
                    except Exception as e:
                        print(f"✗ {file_path.name}: {e}")

        return all_ids

    def search(
        self,
        query: str,
        n_results: int = 5,
        where: Optional[Dict] = None
    ) -> List[Dict[str, Any]]:
        """
        Semantic search in the collection.

        Args:
            query: Search query
            n_results: Number of results to return
            where: Optional filter dict

        Returns:
            List of results with content, metadata, and score
        """
        # Generate query embedding
        query_embedding = self.embedder.encode(query).tolist()

        # Search
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
            include=["documents", "metadatas", "distances"]
        )

        # Format results
        output = []
        for i in range(len(results["ids"][0])):
            output.append({
                "id": results["ids"][0][i],
                "content": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "score": 1 - results["distances"][0][i]  # Convert distance to similarity
            })

        return output

    def query(
        self,
        question: str,
        n_results: int = 5,
        model: str = "qwen2.5:7b",
        system_prompt: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        RAG query - search + LLM analysis.

        Args:
            question: Question to answer
            n_results: Number of documents to retrieve
            model: Ollama model to use
            system_prompt: Optional system prompt

        Returns:
            Dict with answer and sources
        """
        if not HTTPX_AVAILABLE:
            raise ImportError("httpx required for LLM queries: pip install httpx")

        # Search for relevant documents
        results = self.search(question, n_results=n_results)

        if not results:
            return {
                "answer": "No relevant documents found.",
                "sources": [],
                "query": question
            }

        # Build context from results
        context = "\n\n---\n\n".join([
            f"[Source: {r['metadata'].get('source', r['id'])}]\n{r['content']}"
            for r in results
        ])

        # Default system prompt
        if system_prompt is None:
            system_prompt = (
                "You are a helpful assistant. Answer questions based on the provided context. "
                "If the answer is not in the context, say so. Cite your sources."
            )

        # Build prompt
        prompt = f"""Context:
{context}

Question: {question}

Answer based on the context above:"""

        # Call Ollama
        try:
            with httpx.Client(timeout=60.0) as client:
                response = client.post(
                    f"{self.ollama_url}/api/generate",
                    json={
                        "model": model,
                        "prompt": prompt,
                        "system": system_prompt,
                        "stream": False
                    }
                )
                response.raise_for_status()
                answer = response.json().get("response", "")
        except Exception as e:
            answer = f"LLM error: {e}"

        return {
            "answer": answer,
            "sources": results,
            "query": question,
            "model": model
        }

    def count(self) -> int:
        """Return number of documents in collection."""
        return self.collection.count()

    def delete(self, ids: List[str]) -> None:
        """Delete documents by ID."""
        self.collection.delete(ids=ids)

    def clear(self) -> None:
        """Clear all documents from collection."""
        self.client.delete_collection(self.collection_name)
        self.collection = self.client.create_collection(
            name=self.collection_name,
            metadata={"hnsw:space": "cosine"}
        )

    def _chunk_text(
        self,
        text: str,
        chunk_size: int,
        chunk_overlap: int
    ) -> List[str]:
        """Split text into overlapping chunks."""
        chunks = []
        start = 0

        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end]

            # Try to break at sentence boundary
            if end < len(text):
                last_period = chunk.rfind(". ")
                if last_period > chunk_size // 2:
                    chunk = chunk[:last_period + 1]
                    end = start + last_period + 1

            chunks.append(chunk.strip())
            start = end - chunk_overlap

        return [c for c in chunks if c]  # Remove empty chunks

    def _read_pdf(self, file_path: Path) -> str:
        """Read text from PDF file."""
        try:
            from pypdf import PdfReader
            reader = PdfReader(str(file_path))
            text = "\n".join(page.extract_text() or "" for page in reader.pages)
            return text
        except ImportError:
            raise ImportError("PDF support requires: pip install rapid-rag[pdf]")

    def __repr__(self) -> str:
        return f"RapidRAG(collection='{self.collection_name}', docs={self.count()})"
