"""RAG pipeline — embed query, search, assemble context, stream LLM response."""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any

from redhat_oracle.core.config import OracleConfig
from redhat_oracle.core.llm.base import LLMProvider
from redhat_oracle.core.session import ChatSession
from redhat_oracle.core.vectordb.store import VectorStore

logger = logging.getLogger(__name__)


@dataclass
class RAGResult:
    """Result from a RAG query."""

    response: str
    sources: list[dict[str, Any]] = field(default_factory=list)
    search_results: list[dict[str, Any]] = field(default_factory=list)


class RAGPipeline:
    """RAG pipeline: search vector store, build context, query LLM."""

    def __init__(
        self,
        config: OracleConfig,
        store: VectorStore,
        llm: LLMProvider,
    ) -> None:
        self._config = config
        self._store = store
        self._llm = llm

    def _build_context(
        self,
        search_results: list[dict[str, Any]],
        session: ChatSession | None = None,
        max_history: int = 10,
    ) -> str:
        """Build context string from search results and chat history."""
        parts: list[str] = []

        # Add retrieved chunks
        if search_results:
            parts.append("## Retrieved Context\n")
            for i, result in enumerate(search_results, 1):
                meta = result.get("metadata", {})
                url = meta.get("url", "")
                if url:
                    source_ref = url
                else:
                    source_ref = "/".join(
                        v for v in (
                            meta.get("source", "?"),
                            meta.get("org", ""),
                            meta.get("repo", ""),
                            meta.get("path", "?"),
                        ) if v
                    )
                score = result.get("score", 0.0)
                parts.append(f"### Source {i} [{source_ref}] (score: {score:.2f})")
                parts.append(result.get("text", ""))
                parts.append("")

        # Add chat history
        if session and session.messages:
            history = session.get_history(max_history)
            parts.append("## Chat History\n")
            for msg in history:
                role = "User" if msg.role == "user" else "Assistant"
                parts.append(f"**{role}:** {msg.content}\n")

        return "\n".join(parts)

    async def query(
        self,
        question: str,
        session: ChatSession | None = None,
        source: str | None = None,
    ) -> RAGResult:
        """Run a full RAG query: search → context → LLM.

        Args:
            question: The user's question.
            session: Optional chat session for history.
            source: Optional source filter for search.

        Returns:
            RAGResult with response and sources.
        """
        # Search vector store
        search_results = self._store.search(
            query=question,
            source=source,
            top_k=self._config.rag.top_k,
            min_score=self._config.rag.min_score,
        )

        # Build context
        context = self._build_context(search_results, session)

        # Assemble prompt
        prompt = f"{context}\n\n## Current Question\n{question}"

        # Generate response
        response = await self._llm.generate(
            prompt=prompt,
            system=self._config.rag.system_prompt,
        )

        # Extract source references
        sources = [
            {
                "source": r.get("metadata", {}).get("source", ""),
                "org": r.get("metadata", {}).get("org", ""),
                "repo": r.get("metadata", {}).get("repo", ""),
                "path": r.get("metadata", {}).get("path", ""),
                "url": r.get("metadata", {}).get("url", ""),
                "score": r.get("score", 0.0),
            }
            for r in search_results
        ]

        return RAGResult(
            response=response,
            sources=sources,
            search_results=search_results,
        )

    async def stream_query(
        self,
        question: str,
        session: ChatSession | None = None,
        source: str | None = None,
    ) -> tuple[AsyncIterator[str], list[dict[str, Any]]]:
        """Run a RAG query with streaming response.

        Returns:
            Tuple of (token stream, list of sources used).
        """
        # Search vector store
        search_results = self._store.search(
            query=question,
            source=source,
            top_k=self._config.rag.top_k,
            min_score=self._config.rag.min_score,
        )

        # Build context
        context = self._build_context(search_results, session)
        prompt = f"{context}\n\n## Current Question\n{question}"

        # Stream response
        token_stream = self._llm.stream(
            prompt=prompt,
            system=self._config.rag.system_prompt,
        )

        sources = [
            {
                "source": r.get("metadata", {}).get("source", ""),
                "org": r.get("metadata", {}).get("org", ""),
                "repo": r.get("metadata", {}).get("repo", ""),
                "path": r.get("metadata", {}).get("path", ""),
                "url": r.get("metadata", {}).get("url", ""),
                "score": r.get("score", 0.0),
            }
            for r in search_results
        ]

        return token_stream, sources
