import logging

from .base_prompt import BasePrompt

logger = logging.getLogger(__name__)


class RetrieverPrompt(BasePrompt):
    """Handles prompts for generating search queries for tutorial retriever"""

    def default_template(self) -> str:
        """Default template for search query generation"""
        return """
You are an expert at generating search queries to find relevant machine learning tutorials. Given the context below, generate a concise and effective search query that will help find the most relevant tutorials for this task.

### Task Description
{task_description}

### Data Structures
{data_prompt}

### User Instruction
{user_input}

### Previous Error Analysis
{error_analysis}

### Selected Tool/Library
{selected_tool}


Based on the above context, generate a search query that will help find tutorials most relevant to this task. The query should:
1. Include key technical terms and concepts
2. Focus on the main task/problem to solve
3. Be concise but specific

IMPORTANT: Respond ONLY with the search query text. Do not include explanations, quotes, or any other formatting.
"""

    def build(self) -> str:
        """Build a prompt for the LLM to generate a search query."""

        # Format the prompt using the template
        prompt = self.template.format(
            task_description=self.manager.task_description,
            selected_tool=self.manager.selected_tool,
            data_prompt=self.manager.data_prompt,
            user_input=self.manager.user_input,
            error_analysis=self.manager.all_previous_error_prompts,
        )

        self.manager.save_and_log_states(
            content=prompt, save_name="retriever_prompt.txt", per_iteration=True, add_uuid=False
        )

        return prompt

    def parse(self, response: str) -> str:
        """Parse the LLM response to extract the search query."""

        self.manager.save_and_log_states(
            content=response, save_name="retriever_response.txt", per_iteration=True, add_uuid=False
        )

        try:
            # Clean the response - take first line and strip whitespace
            query = response.strip().split("\n")[0].strip()

            # Remove quotes if present
            query = query.strip("\"'")

            # Basic validation
            if not query:
                logger.warning("Empty search query generated by LLM")
                return ""

            # Limit query length for practical purposes
            if len(query) > 512:
                query = query[:512]
                logger.warning("Search query truncated to 512 characters")

            # Remove any unwanted prefixes/suffixes that might indicate explanation
            unwanted_prefixes = [
                "search query:",
                "query:",
                "the search query is:",
            ]
            query_lower = query.lower()
            for prefix in unwanted_prefixes:
                if query_lower.startswith(prefix):
                    query = query[len(prefix) :].strip()
                    break

            logger.info(f"Generated search query: '{query}'")

            self.manager.save_and_log_states(
                content=query, save_name="parsed_search_query.txt", per_iteration=True, add_uuid=False
            )

            return query

        except Exception as e:
            logger.warning(f"Error parsing search query from LLM response: {e}")
            return ""
