import asyncio
import hashlib
from dataclasses import dataclass
from typing import Optional
import asyncpg
from openai import AsyncOpenAI
import tiktoken

@dataclass
class Chunk:
    content: str
    token_count: int
    chunk_index: int
    metadata: dict

class DocumentProcessor:
    def __init__(self, openai_client: AsyncOpenAI, pool: asyncpg.Pool):
        self.openai = openai_client
        self.pool = pool
        self.encoder = tiktoken.encoding_for_model("text-embedding-3-small")
        self.max_tokens = 512
        self.overlap_tokens = 50

    def chunk_text(self, text: str, metadata: dict = {}) -> list[Chunk]:
        """Split text into overlapping chunks with token counting."""
        tokens = self.encoder.encode(text)
        chunks = []
        chunk_index = 0

        start = 0
        while start < len(tokens):
            end = min(start + self.max_tokens, len(tokens))
            chunk_tokens = tokens[start:end]
            chunk_text = self.encoder.decode(chunk_tokens)

            chunks.append(Chunk(
                content=chunk_text,
                token_count=len(chunk_tokens),
                chunk_index=chunk_index,
                metadata=metadata,
            ))

            if end == len(tokens):
                break

            # Overlap: next chunk starts overlap_tokens before current end
            start = end - self.overlap_tokens
            chunk_index += 1

        return chunks

    async def embed_chunks(self, chunks: list[Chunk]) -> list[list[float]]:
        """Batch embed chunks — max 2048 per API call."""
        texts = [c.content for c in chunks]
        batch_size = 100
        all_embeddings = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            response = await self.openai.embeddings.create(
                model="text-embedding-3-small",
                input=batch,
                encoding_format="float",
            )
            embeddings = [item.embedding for item in response.data]
            all_embeddings.extend(embeddings)

        return all_embeddings

    async def ingest_document(
        self,
        title: str,
        content: str,
        source_url: Optional[str] = None,
        metadata: dict = {},
    ) -> str:
        """Full ingestion pipeline: chunk → embed → store."""
        # Create document record
        doc_id = await self.pool.fetchval(
            """INSERT INTO documents (title, content, source_url, metadata)
               VALUES ($1, $2, $3, $4) RETURNING id""",
            title, content, source_url, metadata,
        )

        # Chunk the content
        chunks = self.chunk_text(content, metadata={"doc_title": title, **metadata})
        print(f"Chunked '{title}' into {len(chunks)} chunks")

        # Embed all chunks
        embeddings = await self.embed_chunks(chunks)
        print(f"Embedded {len(embeddings)} chunks")

        # Bulk insert chunks + embeddings
        async with self.pool.acquire() as conn:
            await conn.executemany(
                """INSERT INTO document_chunks
                   (document_id, chunk_index, content, token_count, embedding, metadata)
                   VALUES ($1, $2, $3, $4, $5::vector, $6)""",
                [
                    (doc_id, chunk.chunk_index, chunk.content, chunk.token_count,
                     str(embedding), chunk.metadata)
                    for chunk, embedding in zip(chunks, embeddings)
                ]
            )

        print(f"Stored document '{title}' (id: {doc_id})")
        return str(doc_id)

from anthropic import AsyncAnthropic

class RAGQueryEngine:
    def __init__(self, openai_client: AsyncOpenAI, anthropic_client: AsyncAnthropic, pool: asyncpg.Pool):
        self.openai = openai_client
        self.anthropic = anthropic_client
        self.pool = pool

    async def embed_query(self, query: str) -> list[float]:
        response = await self.openai.embeddings.create(
            model="text-embedding-3-small",
            input=query,
        )
        return response.data[0].embedding

    async def semantic_search(
        self,
        query: str,
        k: int = 10,
        similarity_threshold: float = 0.7,
        metadata_filter: Optional[dict] = None,
    ) -> list[dict]:
        """Retrieve top-k semantically similar chunks."""
        query_embedding = await self.embed_query(query)

        # Build filter clause
        filter_clause = ""
        params = [str(query_embedding), k, similarity_threshold]

        if metadata_filter:
            filter_clause = "AND metadata @> $4::jsonb"
            params.append(str(metadata_filter))

        results = await self.pool.fetch(
            f"""SELECT
                dc.id,
                dc.content,
                dc.metadata,
                d.title,
                d.source_url,
                1 - (dc.embedding <=> $1::vector) AS similarity
               FROM document_chunks dc
               JOIN documents d ON d.id = dc.document_id
               WHERE 1 - (dc.embedding <=> $1::vector) > $3
               {filter_clause}
               ORDER BY dc.embedding <=> $1::vector
               LIMIT $2""",
            *params,
        )

        return [dict(r) for r in results]

    async def rag_query(
        self,
        question: str,
        k: int = 8,
        system_context: str = "",
    ) -> dict:
        """Full RAG: retrieve context + generate answer."""
        # 1. Retrieve relevant chunks
        chunks = await self.semantic_search(question, k=k)

        if not chunks:
            return {
                "answer": "I couldn't find relevant information to answer this question.",
                "sources": [],
                "context_used": 0,
            }

        # 2. Build context from retrieved chunks (with deduplication)
        seen_docs = set()
        context_parts = []
        sources = []

        for chunk in chunks:
            context_parts.append(
                f"[Source: {chunk['title']}]\n{chunk['content']}"
            )
            if chunk['title'] not in seen_docs:
                seen_docs.add(chunk['title'])
                sources.append({
                    "title": chunk['title'],
                    "url": chunk.get('source_url'),
                    "similarity": round(chunk['similarity'], 3),
                })

        context = "\n\n---\n\n".join(context_parts)

        # 3. Generate answer with context
        system_prompt = f"""You are a knowledgeable assistant. Answer the user's question based ONLY on the provided context.
If the context doesn't contain enough information, say so clearly.
Do not make up information not present in the context.

{system_context}

Context:
{context}"""

        response = await self.anthropic.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=1024,
            system=system_prompt,
            messages=[{"role": "user", "content": question}],
        )

        return {
            "answer": response.content[0].text,
            "sources": sources,
            "context_used": len(chunks),
            "input_tokens": response.usage.input_tokens,
            "output_tokens": response.usage.output_tokens,
        }

async def main():
    pool = await asyncpg.create_pool("postgresql://localhost/vector_demo")
    openai = AsyncOpenAI()
    anthropic = AsyncAnthropic()

    processor = DocumentProcessor(openai, pool)
    engine = RAGQueryEngine(openai, anthropic, pool)

    # Ingest documents
    await processor.ingest_document(
        title="Q4 2025 Product Roadmap",
        content="Our Q4 priorities are: 1) Launch the new payments API...",
        metadata={"category": "product", "quarter": "Q4-2025"},
    )

    # Query
    result = await engine.rag_query(
        "What are the Q4 2025 payment priorities?"
    )

    print(f"Answer: {result['answer']}")
    print(f"Sources: {[s['title'] for s in result['sources']]}")
    print(f"Context chunks used: {result['context_used']}")

async def hybrid_search(self, query: str, k: int = 10, alpha: float = 0.7) -> list[dict]:
    """
    alpha=1.0: pure vector search
    alpha=0.0: pure BM25 (keyword)
    alpha=0.7: 70% vector, 30% BM25 (recommended default)
    """
    # Vector search results
    vector_results = await self.semantic_search(query, k=k*2)

    # BM25 search via PostgreSQL full-text search
    bm25_results = await self.pool.fetch(
        """SELECT dc.id, dc.content, d.title,
                  ts_rank(to_tsvector('english', dc.content),
                          plainto_tsquery('english', $1)) AS bm25_score
           FROM document_chunks dc
           JOIN documents d ON d.id = dc.document_id
           WHERE to_tsvector('english', dc.content) @@ plainto_tsquery('english', $1)
           ORDER BY bm25_score DESC
           LIMIT $2""",
        query, k * 2,
    )

    # Reciprocal Rank Fusion (RRF)
    scores = {}
    for rank, chunk in enumerate(vector_results):
        scores[chunk['id']] = {'data': chunk, 'score': alpha * (1 / (rank + 60))}

    for rank, chunk in enumerate(bm25_results):
        chunk_id = str(chunk['id'])
        if chunk_id in scores:
            scores[chunk_id]['score'] += (1 - alpha) * (1 / (rank + 60))
        else:
            scores[chunk_id] = {'data': dict(chunk), 'score': (1 - alpha) * (1 / (rank + 60))}

    return sorted(scores.values(), key=lambda x: x['score'], reverse=True)[:k]

import hashlib
import redis.asyncio as redis

class CachedEmbedder:
    def __init__(self, openai_client: AsyncOpenAI, redis_client: redis.Redis):
        self.openai = openai_client
        self.redis = redis_client
        self.ttl = 86400 * 30  # 30 days

    def _cache_key(self, text: str, model: str) -> str:
        hash_val = hashlib.sha256(f"{model}:{text}".encode()).hexdigest()
        return f"embedding:{hash_val}"

    async def embed(self, text: str, model: str = "text-embedding-3-small") -> list[float]:
        key = self._cache_key(text, model)

        # Check cache
        cached = await self.redis.get(key)
        if cached:
            import json
            return json.loads(cached)

        # Generate and cache
        response = await self.openai.embeddings.create(model=model, input=text)
        embedding = response.data[0].embedding

        await self.redis.setex(key, self.ttl, __import__('json').dumps(embedding))
        return embedding