"""
RAG decorator for LangChain/LangGraph agents.

Automatically injects knowledge base context into any agent.

Usage:
    from ragsaas import with_rag

    # Wrap an existing agent
    @with_rag(tenant_id="customer-123")
    agent = create_react_agent(llm, tools)

    # Or use with dynamic tenant (e.g., from request context)
    @with_rag(tenant_id=lambda: get_current_user().tenant_id)
    async def handle_query(query: str):
        return await agent.invoke(query)
"""

import os
import functools
from typing import Optional, Callable, Union, Any, List
from dataclasses import dataclass
import httpx


@dataclass
class RAGContext:
    """Context retrieved from knowledge base."""
    content: str
    sources: List[dict]


class RagSaaSMiddleware:
    """Middleware that injects RAG context into agent calls."""

    def __init__(
        self,
        tenant_id: Union[str, Callable[[], str]],
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        collection: Optional[str] = None,
        top_k: int = 5,
        use_hybrid: bool = False,
        auto_inject: bool = True,
        threshold: Optional[float] = None,
        rewrite_query: bool = False,
        rewrite_model: str = "gpt-4o-mini",
    ):
        """
        Initialize RAG middleware.

        Args:
            tenant_id: Tenant identifier (string or callable that returns string)
            api_key: API key (defaults to RAGSAAS_API_KEY env var)
            base_url: API base URL (defaults to RAGSAAS_BASE_URL or https://api.ragsaas.com)
            collection: Filter by collection name
            top_k: Number of results to retrieve
            use_hybrid: Use hybrid search (semantic + keyword)
            auto_inject: Automatically inject context into messages
            threshold: Minimum score threshold
            rewrite_query: Use LLM to rewrite ambiguous queries (adds latency but more accurate)
            rewrite_model: Model to use for query rewriting (default: gpt-4o-mini)
        """
        self._tenant_id = tenant_id
        self.api_key = api_key or os.environ.get("RAGSAAS_API_KEY")
        self.base_url = (base_url or os.environ.get("RAGSAAS_BASE_URL", "https://api.ragsaas.com")).rstrip("/")
        self.collection = collection
        self.top_k = top_k
        self.use_hybrid = use_hybrid
        self.auto_inject = auto_inject
        self.threshold = threshold
        self.rewrite_query = rewrite_query
        self.rewrite_model = rewrite_model

        if not self.api_key:
            raise ValueError(
                "RAGSAAS_API_KEY environment variable is required. "
                "Get your API key at https://ragsaas.com/dashboard"
            )

    @property
    def tenant_id(self) -> str:
        """Get tenant ID (resolves callable if needed)."""
        if callable(self._tenant_id):
            return self._tenant_id()
        return self._tenant_id

    def _get_client(self) -> httpx.Client:
        """Get HTTP client."""
        return httpx.Client(
            base_url=self.base_url,
            headers={"X-API-Key": self.api_key},
            timeout=30.0,
        )

    def retrieve(self, query: str) -> RAGContext:
        """Retrieve relevant context for a query."""
        endpoint = "/v1/query/hybrid" if self.use_hybrid else "/v1/query"

        payload = {
            "tenant_id": self.tenant_id,
            "query": query,
            "top_k": self.top_k,
        }
        if self.collection:
            payload["collection"] = self.collection
        if self.threshold is not None:
            payload["threshold"] = self.threshold
        if self.use_hybrid:
            payload["semantic_weight"] = 0.7

        with self._get_client() as client:
            response = client.post(endpoint, json=payload)
            response.raise_for_status()
            data = response.json()

        # Format context
        results = data.get("results", [])
        if not results:
            return RAGContext(content="", sources=[])

        context_parts = []
        sources = []
        for i, r in enumerate(results):
            context_parts.append(f"[{i+1}] From '{r['document_name']}':\n{r['content']}")
            sources.append({
                "index": i + 1,
                "document_name": r["document_name"],
                "document_id": r["document_id"],
                "score": r["score"],
            })

        return RAGContext(
            content="\n\n".join(context_parts),
            sources=sources,
        )

    def rewrite_query_with_context(self, query: str, history: List[dict]) -> str:
        """Use LLM to rewrite an ambiguous query with conversation context.

        Args:
            query: The original user query
            history: Recent conversation history

        Returns:
            Rewritten self-contained query
        """
        if not history:
            return query

        # Format history
        history_text = []
        for msg in history[-6:]:  # Last 3 turns
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role and content:
                history_text.append(f"{role}: {content[:200]}")

        prompt = f"""Given this conversation history:
{chr(10).join(history_text)}

The user's latest question is: "{query}"

Rewrite this question to be self-contained and clear, incorporating any necessary context from the conversation.
Only respond with the rewritten question, nothing else."""

        try:
            # Use OpenAI directly for rewriting
            import openai
            client = openai.OpenAI()
            response = client.chat.completions.create(
                model=self.rewrite_model,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=150,
                temperature=0,
            )
            return response.choices[0].message.content.strip()
        except Exception:
            # Fallback to original query if rewriting fails
            return query

    def build_system_prompt(self, context: RAGContext, original_system: str = "") -> str:
        """Build system prompt with RAG context."""
        if not context.content:
            return original_system

        rag_prompt = f"""You have access to a knowledge base. Use the following context to answer questions.
Cite sources using [1], [2], etc. If the answer is not in the context, say so.

Knowledge Base Context:
{context.content}"""

        if original_system:
            return f"{rag_prompt}\n\n---\n\nAdditional Instructions:\n{original_system}"
        return rag_prompt

    def wrap_messages(self, messages: List[dict], query: str) -> List[dict]:
        """Wrap messages with RAG context.

        Inserts RAG context before the last user message so the static
        prefix (system prompt + conversation history) remains cacheable
        by LLM providers.
        """
        context = self.retrieve(query)

        if not context.content:
            return messages

        messages = list(messages)  # Copy

        rag_message = {
            "role": "system",
            "content": self.build_system_prompt(context),
        }

        # Find last user message and insert RAG context right before it
        last_user_idx = None
        for i in range(len(messages) - 1, -1, -1):
            if messages[i].get("role") == "user":
                last_user_idx = i
                break

        if last_user_idx is not None:
            messages.insert(last_user_idx, rag_message)
        else:
            messages.append(rag_message)

        return messages


def with_rag(
    tenant_id: Union[str, Callable[[], str]],
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    collection: Optional[str] = None,
    top_k: int = 5,
    use_hybrid: bool = False,
    threshold: Optional[float] = None,
    rewrite_query: bool = False,
    rewrite_model: str = "gpt-4o-mini",
):
    """
    Decorator to add RAG capabilities to a LangChain/LangGraph agent.

    Usage:
        # With a ReAct agent
        @with_rag(tenant_id="customer-123")
        def create_agent():
            return create_react_agent(llm, tools)

        agent = create_agent()
        result = agent.invoke({"messages": [("user", "What is X?")]})

        # With dynamic tenant
        @with_rag(tenant_id=lambda: request.user.tenant_id)
        async def query_handler(query: str):
            return await agent.ainvoke({"messages": [("user", query)]})

        # With query rewriting for better accuracy
        @with_rag(tenant_id="customer-123", rewrite_query=True)
        def create_agent():
            return create_react_agent(llm, tools)

    Args:
        tenant_id: Tenant identifier (string or callable)
        api_key: API key (defaults to RAGSAAS_API_KEY env var)
        base_url: API base URL
        collection: Filter by collection name
        top_k: Number of results to retrieve
        use_hybrid: Use hybrid search
        threshold: Minimum score threshold
        rewrite_query: Use LLM to rewrite ambiguous queries (more accurate, adds ~100ms)
        rewrite_model: Model for query rewriting (default: gpt-4o-mini)
    """
    middleware = RagSaaSMiddleware(
        tenant_id=tenant_id,
        api_key=api_key,
        base_url=base_url,
        collection=collection,
        top_k=top_k,
        use_hybrid=use_hybrid,
        threshold=threshold,
        rewrite_query=rewrite_query,
        rewrite_model=rewrite_model,
    )

    def decorator(func_or_class):
        # Check if it's a LangGraph agent (has invoke method)
        if hasattr(func_or_class, 'invoke'):
            return _wrap_agent(func_or_class, middleware)

        # Otherwise wrap the function
        @functools.wraps(func_or_class)
        def wrapper(*args, **kwargs):
            result = func_or_class(*args, **kwargs)
            # If result is an agent, wrap it
            if hasattr(result, 'invoke'):
                return _wrap_agent(result, middleware)
            return result

        return wrapper

    return decorator


def _wrap_agent(agent: Any, middleware: RagSaaSMiddleware) -> Any:
    """Wrap a LangChain/LangGraph agent with RAG middleware."""

    class RAGWrappedAgent:
        """Agent wrapper that injects RAG context."""

        # Markers that indicate query needs context
        INCOMPLETE_MARKERS = [
            # Spanish
            "esto", "eso", "aquello", "el mismo", "la misma",
            "este", "ese", "y el", "y la", "y los", "y las",
            "qué más", "algo más", "también", "además",
            "cuánto cuesta", "cuál es el precio", "lo mismo",
            # English
            "this", "that", "it", "the same", "these", "those",
            "what about", "how much", "and the", "also", "too",
        ]

        def __init__(self, agent, middleware):
            self._agent = agent
            self._middleware = middleware
            # Copy attributes from original agent
            for attr in dir(agent):
                if not attr.startswith('_') and attr not in ('invoke', 'ainvoke', 'stream', 'astream'):
                    try:
                        setattr(self, attr, getattr(agent, attr))
                    except AttributeError:
                        pass

        def _parse_message(self, msg) -> tuple:
            """Parse message to (role, content) tuple."""
            if isinstance(msg, tuple):
                return msg[0], msg[1]
            elif isinstance(msg, dict):
                return msg.get("role", ""), msg.get("content", "")
            elif hasattr(msg, 'type'):
                role = "assistant" if msg.type == "ai" else "user" if msg.type == "human" else msg.type
                return role, getattr(msg, 'content', '')
            return "", ""

        def _is_complete_query(self, query: str) -> bool:
            """Check if query is self-contained or needs history context."""
            if not query:
                return False

            query_lower = query.lower().strip()

            # Very short queries likely need context
            if len(query.split()) < 3:
                return False

            # Check for incomplete markers
            for marker in self.INCOMPLETE_MARKERS:
                if marker in query_lower:
                    return False

            # Starts with conjunction
            if query_lower.startswith(("y ", "and ", "pero ", "but ", "o ", "or ")):
                return False

            return True

        def _build_search_query(self, input_data: dict) -> str:
            """Build search query, including history context if needed."""
            messages = input_data.get("messages", [])
            if not messages:
                return ""

            # Get last user message
            last_query = ""
            for msg in reversed(messages):
                role, content = self._parse_message(msg)
                if role == "user":
                    last_query = content
                    break

            if not last_query:
                return ""

            # If query is complete, use it directly
            if self._is_complete_query(last_query):
                return last_query

            # Otherwise, include recent history for context
            context_parts = []
            for msg in messages[-6:]:  # Last 3 turns max
                role, content = self._parse_message(msg)
                if role in ("user", "assistant") and content:
                    # Truncate long messages
                    truncated = content[:300] + "..." if len(content) > 300 else content
                    context_parts.append(f"{role}: {truncated}")

            return "\n".join(context_parts)

        def _extract_query(self, input_data: dict) -> str:
            """Extract user query from input (legacy method, use _build_search_query)."""
            messages = input_data.get("messages", [])

            # Find last user message
            for msg in reversed(messages):
                if isinstance(msg, tuple) and msg[0] == "user":
                    return msg[1]
                elif isinstance(msg, dict) and msg.get("role") == "user":
                    return msg.get("content", "")
                elif hasattr(msg, 'type') and msg.type == "human":
                    return msg.content

            return ""

        def _inject_context(self, input_data: dict) -> dict:
            """Inject RAG context into input."""
            messages = input_data.get("messages", [])

            # If LLM rewriting is enabled, use it for ambiguous queries
            if self._middleware.rewrite_query:
                last_query = self._extract_query(input_data)
                if last_query and not self._is_complete_query(last_query):
                    # Convert messages to dict format for rewriting
                    history = []
                    for msg in messages[-6:]:
                        role, content = self._parse_message(msg)
                        if role and content:
                            history.append({"role": role, "content": content})
                    search_query = self._middleware.rewrite_query_with_context(last_query, history)
                else:
                    search_query = last_query
            else:
                # Use simple concatenation approach
                search_query = self._build_search_query(input_data)

            if not search_query:
                return input_data

            context = self._middleware.retrieve(search_query)
            if not context.content:
                return input_data

            # Create a copy of input
            input_data = dict(input_data)
            messages = list(input_data.get("messages", []))

            # Build RAG context message.
            # Insert BEFORE the last user message (not at the beginning)
            # so the static prefix (system prompt + history) stays cacheable
            # by LLM providers (Anthropic prompt caching, OpenAI prefix caching).
            rag_content = self._middleware.build_system_prompt(context)

            try:
                from langchain_core.messages import SystemMessage
                rag_message = SystemMessage(content=rag_content)
            except ImportError:
                rag_message = {"role": "system", "content": rag_content}

            # Find the last user message and insert RAG context right before it
            last_user_idx = None
            for i in range(len(messages) - 1, -1, -1):
                role, _ = self._parse_message(messages[i])
                if role == "user":
                    last_user_idx = i
                    break

            if last_user_idx is not None:
                messages.insert(last_user_idx, rag_message)
            else:
                messages.append(rag_message)

            input_data["messages"] = messages

            return input_data

        def invoke(self, input_data: dict, **kwargs):
            """Invoke agent with RAG context."""
            enhanced_input = self._inject_context(input_data)
            return self._agent.invoke(enhanced_input, **kwargs)

        async def ainvoke(self, input_data: dict, **kwargs):
            """Async invoke agent with RAG context."""
            enhanced_input = self._inject_context(input_data)
            return await self._agent.ainvoke(enhanced_input, **kwargs)

        def stream(self, input_data: dict, **kwargs):
            """Stream agent with RAG context."""
            enhanced_input = self._inject_context(input_data)
            return self._agent.stream(enhanced_input, **kwargs)

        async def astream(self, input_data: dict, **kwargs):
            """Async stream agent with RAG context."""
            enhanced_input = self._inject_context(input_data)
            return self._agent.astream(enhanced_input, **kwargs)

        def __getattr__(self, name):
            return getattr(self._agent, name)

    return RAGWrappedAgent(agent, middleware)


# Convenience class for manual usage
class RAGAgent:
    """
    Convenience class to create a RAG-enhanced agent.

    Usage:
        rag = RAGAgent(tenant_id="customer-123")
        agent = rag.wrap(create_react_agent(llm, tools))

        # Or retrieve context manually
        context = rag.retrieve("What is X?")
    """

    def __init__(
        self,
        tenant_id: Union[str, Callable[[], str]],
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        collection: Optional[str] = None,
        top_k: int = 5,
        use_hybrid: bool = False,
        threshold: Optional[float] = None,
    ):
        self._middleware = RagSaaSMiddleware(
            tenant_id=tenant_id,
            api_key=api_key,
            base_url=base_url,
            collection=collection,
            top_k=top_k,
            use_hybrid=use_hybrid,
            threshold=threshold,
        )

    def retrieve(self, query: str) -> RAGContext:
        """Retrieve relevant context for a query."""
        return self._middleware.retrieve(query)

    def wrap(self, agent: Any) -> Any:
        """Wrap an agent with RAG capabilities."""
        return _wrap_agent(agent, self._middleware)
