# Copyright (c) 2025 Microsoft Corporation.
# Licensed under the MIT License

"""LLM-based query rewriting for keyword-based retrieval.

Platform search APIs work best with short keyword queries rather than natural
language questions. This module uses the configured LLM to rewrite complex
queries into 1-5 simpler keyword queries.
"""

import logging

from nlweb_core.llm import ask_llm
from pydantic import BaseModel, Field, field_validator

logger = logging.getLogger(__name__)


class QueryRewriteResponse(BaseModel):
    """Response model for query rewriting into keyword queries."""

    rewritten_queries: list[str] = Field(
        ...,
        min_length=1,
        max_length=5,
        description="List of 1-5 simpler keyword queries",
    )

    @field_validator("rewritten_queries", mode="before")
    @classmethod
    def filter_empty_queries(cls, v: list[str]) -> list[str]:
        """Filter out empty queries and limit to 5."""
        if isinstance(v, list):
            filtered = [q.strip() for q in v if q and isinstance(q, str) and q.strip()]
            return filtered[:5]
        return v


QUERY_REWRITE_PROMPT = """You are helping to rewrite a complex search query into simpler keyword queries for a traditional keyword-based search engine.
The search engine works best with short, focused queries containing important keywords.

Take the following query and break it down into up to 5 simpler search queries.
Each query should:
- Contain no more than 3 words
- Focus on the most important keywords and concepts
- Be diverse to cover different aspects of the original query
- Use only essential nouns, adjectives, or product terms
- Avoid common words like "for", "the", "some", "are", "that", "would", "be"

For example:
- "what are some options for plates that would be appropriate for serving vegetables" → ["vegetable plates", "serving plates", "dinner plates", "salad plates", "ceramic plates"]
- "looking for a tea pot that can brew green tea" → ["tea pot", "green tea", "teapot ceramic", "japanese teapot", "brewing pot"]

The original query is: {query}"""


async def rewrite_query(query: str) -> list[str]:
    """Rewrite a complex query into simpler keyword queries.

    Uses the configured 'high' LLM provider to break down the query.

    Args:
        query: The original natural language query.

    Returns:
        List of 1-5 keyword queries. Falls back to [query] on failure.
    """
    try:
        prompt = QUERY_REWRITE_PROMPT.format(query=query)
        response = await ask_llm(
            prompt=prompt,
            schema=QueryRewriteResponse,
            level="high",
            timeout=10,
            max_length=512,
        )
        if response.rewritten_queries:
            return response.rewritten_queries
    except Exception as e:
        logger.warning(f"Query rewrite failed, using original query: {e}")

    return [query]
