"""LLM client for intelligent code search using OpenAI, OpenRouter, or AWS Bedrock API."""

import asyncio
import json
import os
import re
from collections.abc import AsyncIterator
from typing import Any, Literal

import httpx
from loguru import logger

from .exceptions import SearchError

# Type alias for provider
LLMProvider = Literal["openai", "openrouter", "bedrock"]

# Type alias for intent
IntentType = Literal["find", "answer", "analyze"]


class LLMClient:
    """Client for LLM-powered intelligent search orchestration.

    Supports OpenAI, OpenRouter, and AWS Bedrock APIs:
    1. Generate multiple targeted search queries from natural language
    2. Analyze search results and select most relevant ones
    3. Provide contextual explanations for results

    Provider Selection Priority:
    1. Explicit provider parameter
    2. Preferred provider from config
    3. Auto-detect: Bedrock (if AWS creds) → OpenRouter → OpenAI

    Default Models:
    - Bedrock: Claude Sonnet 4.6 (us.anthropic.claude-sonnet-4-6-20250514-v1:0)
    - OpenRouter: Claude Opus 4.5
    - OpenAI: GPT-4o-mini
    """

    # Default models for each provider (comparable performance/cost)
    DEFAULT_MODELS = {
        "openai": "gpt-4o-mini",  # Fast, cheap, comparable to claude-3-haiku
        "openrouter": "anthropic/claude-opus-4.5",  # Claude Opus 4.5 for chat REPL
        "bedrock": "us.anthropic.claude-sonnet-4-6-20250514-v1:0",  # Claude Sonnet 4.6 (cross-region inference)
    }

    # Advanced "thinking" models for complex queries (--think flag)
    THINKING_MODELS = {
        "openai": "gpt-4o",  # More capable, better reasoning
        "openrouter": "anthropic/claude-opus-4.5",  # Claude Opus 4.5 for deep analysis
        "bedrock": "anthropic.claude-opus-4-20250514-v1:0",  # Claude Opus 4
    }

    # API endpoints
    API_ENDPOINTS = {
        "openai": "https://api.openai.com/v1/chat/completions",
        "openrouter": "https://openrouter.ai/api/v1/chat/completions",
    }

    TIMEOUT_SECONDS = 30.0

    def __init__(
        self,
        api_key: str | None = None,
        model: str | None = None,
        timeout: float = TIMEOUT_SECONDS,
        provider: LLMProvider | None = None,
        openai_api_key: str | None = None,
        openrouter_api_key: str | None = None,
        think: bool = False,
    ) -> None:
        """Initialize LLM client.

        Args:
            api_key: API key (deprecated, use provider-specific keys)
            model: Model to use (defaults based on provider)
            timeout: Request timeout in seconds
            provider: Explicit provider ('openai', 'openrouter', or 'bedrock')
            openai_api_key: OpenAI API key (or use OPENAI_API_KEY env var)
            openrouter_api_key: OpenRouter API key (or use OPENROUTER_API_KEY env var)
            think: Use advanced "thinking" model for complex queries

        Raises:
            ValueError: If no API key/credentials found for any provider
        """
        self.think = think
        self._bedrock_client = None  # Lazy initialization

        # Get API keys from environment or parameters
        self.openai_key = openai_api_key or os.environ.get("OPENAI_API_KEY")
        self.openrouter_key = openrouter_api_key or os.environ.get("OPENROUTER_API_KEY")

        # Support deprecated api_key parameter (assume OpenRouter for backward compatibility)
        if api_key and not self.openrouter_key:
            self.openrouter_key = api_key

        # Determine which provider to use
        if provider:
            # Explicit provider specified
            self.provider: LLMProvider = provider
            if provider == "openai" and not self.openai_key:
                raise ValueError(
                    "OpenAI provider specified but OPENAI_API_KEY not found. "
                    "Please set OPENAI_API_KEY environment variable."
                )
            elif provider == "openrouter" and not self.openrouter_key:
                raise ValueError(
                    "OpenRouter provider specified but OPENROUTER_API_KEY not found. "
                    "Please set OPENROUTER_API_KEY environment variable."
                )
            elif provider == "bedrock" and not self._bedrock_available:
                raise ValueError(
                    "Bedrock provider specified but AWS credentials not found. "
                    "Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables."
                )
        else:
            # Auto-detect provider (prefer Bedrock → OpenRouter → OpenAI)
            if self._bedrock_available:
                self.provider = "bedrock"
            elif self.openrouter_key:
                self.provider = "openrouter"
            elif self.openai_key:
                self.provider = "openai"
            else:
                raise ValueError(
                    "No API key or AWS credentials found. Please set AWS credentials "
                    "(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) for Bedrock, "
                    "OPENROUTER_API_KEY for OpenRouter, or OPENAI_API_KEY for OpenAI."
                )

        # Set API key and endpoint based on provider
        # Select model: explicit > env var > thinking model > default model
        if self.provider == "openai":
            self.api_key = self.openai_key
            self.api_endpoint = self.API_ENDPOINTS["openai"]
            default_model = (
                self.THINKING_MODELS["openai"]
                if think
                else self.DEFAULT_MODELS["openai"]
            )
            self.model = model or os.environ.get("OPENAI_MODEL", default_model)
        elif self.provider == "openrouter":
            self.api_key = self.openrouter_key
            self.api_endpoint = self.API_ENDPOINTS["openrouter"]
            default_model = (
                self.THINKING_MODELS["openrouter"]
                if think
                else self.DEFAULT_MODELS["openrouter"]
            )
            self.model = model or os.environ.get("OPENROUTER_MODEL", default_model)
        else:  # bedrock
            self.api_key = None  # Not used for Bedrock
            self.api_endpoint = None  # Not used for Bedrock
            default_model = (
                self.THINKING_MODELS["bedrock"]
                if think
                else self.DEFAULT_MODELS["bedrock"]
            )
            self.model = model or os.environ.get("BEDROCK_MODEL", default_model)

        self.timeout = timeout

        logger.debug(
            f"Initialized LLM client with provider: {self.provider}, model: {self.model}"
        )

    @property
    def _bedrock_available(self) -> bool:
        """Check if AWS Bedrock credentials are available.

        Checks for:
        1. Explicit AWS credentials (AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY)
        2. AWS session token (AWS_SESSION_TOKEN) - used with temporary credentials
        3. Custom Bedrock endpoint (ANTHROPIC_BEDROCK_BASE_URL) - Claude Code compatibility

        Note: boto3 also supports ~/.aws/credentials and instance profiles,
        but we can't easily check those without importing boto3.
        """
        # Check explicit credentials
        has_explicit_creds = bool(
            os.environ.get("AWS_ACCESS_KEY_ID")
            and os.environ.get("AWS_SECRET_ACCESS_KEY")
        )

        # Check for custom Bedrock endpoint (Claude Code sets this)
        has_bedrock_endpoint = bool(os.environ.get("ANTHROPIC_BEDROCK_BASE_URL"))

        return has_explicit_creds or has_bedrock_endpoint

    def _get_bedrock_client(self) -> Any:
        """Get or create boto3 bedrock-runtime client (lazy initialization).

        Supports:
        - Standard boto3 credential chain (env vars, ~/.aws/credentials, instance profiles)
        - AWS_REGION / AWS_DEFAULT_REGION env vars (defaults to us-east-1)
        - ANTHROPIC_BEDROCK_BASE_URL for custom endpoint (Claude Code compatibility)

        Returns:
            boto3 bedrock-runtime client

        Raises:
            ImportError: If boto3 is not installed
        """
        if self._bedrock_client is None:
            try:
                import boto3
            except ImportError as e:
                raise ImportError(
                    "boto3 is required for Bedrock support. Install with: pip install boto3"
                ) from e

            # Get AWS region from environment or use default
            region = os.environ.get("AWS_REGION") or os.environ.get(
                "AWS_DEFAULT_REGION", "us-east-1"
            )

            # Check for custom Bedrock endpoint (Claude Code sets this)
            bedrock_endpoint = os.environ.get("ANTHROPIC_BEDROCK_BASE_URL")

            if bedrock_endpoint:
                # Custom endpoint specified (Claude Code or custom setup)
                self._bedrock_client = boto3.client(
                    service_name="bedrock-runtime",
                    region_name=region,
                    endpoint_url=bedrock_endpoint,
                )
                logger.debug(
                    f"Initialized Bedrock client in region: {region} with custom endpoint: {bedrock_endpoint}"
                )
            else:
                # Standard Bedrock endpoint
                self._bedrock_client = boto3.client(
                    service_name="bedrock-runtime",
                    region_name=region,
                )
                logger.debug(f"Initialized Bedrock client in region: {region}")

        return self._bedrock_client

    async def generate_search_queries(
        self, natural_language_query: str, limit: int = 3
    ) -> list[str]:
        """Generate targeted search queries from natural language.

        Args:
            natural_language_query: User's natural language query
            limit: Maximum number of search queries to generate

        Returns:
            List of targeted search queries

        Raises:
            SearchError: If API call fails
        """
        system_prompt = """You are a code search expert. Your task is to convert natural language questions about code into targeted search queries.

Given a natural language query, generate {limit} specific search queries that will help find the relevant code.

Rules:
1. Each query should target a different aspect of the question
2. Use technical terms and identifiers when possible
3. Keep queries concise (3-7 words each)
4. Focus on code patterns, function names, class names, or concepts
5. Return ONLY the search queries, one per line, no explanations

Example:
Input: "where is the similarity_threshold parameter set?"
Output:
similarity_threshold default value
similarity_threshold configuration
SemanticSearchEngine init threshold"""

        user_prompt = f"""Natural language query: {natural_language_query}

Generate {limit} targeted search queries:"""

        try:
            messages = [
                {"role": "system", "content": system_prompt.format(limit=limit)},
                {"role": "user", "content": user_prompt},
            ]

            response = await self._chat_completion(messages)

            # Parse queries from response
            content = (
                response.get("choices", [{}])[0].get("message", {}).get("content", "")
            )
            queries = [q.strip() for q in content.strip().split("\n") if q.strip()]

            logger.debug(
                f"Generated {len(queries)} search queries from: '{natural_language_query}'"
            )

            return queries[:limit]

        except Exception as e:
            logger.error(f"Failed to generate search queries: {e}")
            raise SearchError(f"LLM query generation failed: {e}") from e

    async def analyze_and_rank_results(
        self,
        original_query: str,
        search_results: dict[str, list[Any]],
        top_n: int = 5,
    ) -> list[dict[str, Any]]:
        """Analyze search results and select the most relevant ones.

        Args:
            original_query: Original natural language query
            search_results: Dictionary mapping search queries to their results
            top_n: Number of top results to return

        Returns:
            List of ranked results with explanations

        Raises:
            SearchError: If API call fails
        """
        # Format results for LLM analysis
        results_summary = self._format_results_for_analysis(search_results)

        system_prompt = """You are a code search expert. Your task is to analyze search results and identify the most relevant ones for answering a user's question.

Given:
1. A natural language query
2. Multiple search results from different queries

Select the top {top_n} most relevant results that best answer the user's question.

For each selected result, provide:
1. Result identifier (e.g., "Query 1, Result 2")
2. Relevance level: "High", "Medium", or "Low"
3. Brief explanation (1-2 sentences) of why this result is relevant

Format your response as:
RESULT: [identifier]
RELEVANCE: [level]
EXPLANATION: [why this matches]

---

Only include the top {top_n} results."""

        user_prompt = f"""Original Question: {original_query}

Search Results:
{results_summary}

Select the top {top_n} most relevant results:"""

        try:
            messages = [
                {"role": "system", "content": system_prompt.format(top_n=top_n)},
                {"role": "user", "content": user_prompt},
            ]

            response = await self._chat_completion(messages)

            # Parse LLM response
            content = (
                response.get("choices", [{}])[0].get("message", {}).get("content", "")
            )

            ranked_results = self._parse_ranking_response(
                content, search_results, top_n
            )

            logger.debug(f"Ranked {len(ranked_results)} results from LLM analysis")

            return ranked_results

        except Exception as e:
            logger.error(f"Failed to analyze results: {e}")
            raise SearchError(f"LLM analysis failed: {e}") from e

    async def _chat_completion(self, messages: list[dict[str, str]]) -> dict[str, Any]:
        """Make chat completion request to OpenAI, OpenRouter, or Bedrock API.

        Args:
            messages: List of message dictionaries with role and content

        Returns:
            API response dictionary

        Raises:
            SearchError: If API request fails
        """
        # Route to Bedrock if that's the provider
        if self.provider == "bedrock":
            return await self._bedrock_chat_completion(messages)

        # Build headers based on provider
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        # OpenRouter-specific headers
        if self.provider == "openrouter":
            headers["HTTP-Referer"] = "https://github.com/bobmatnyc/mcp-vector-search"
            headers["X-Title"] = "MCP Vector Search"

        payload = {
            "model": self.model,
            "messages": messages,
        }

        provider_name = self.provider.capitalize()

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    self.api_endpoint,
                    headers=headers,
                    json=payload,
                )

                response.raise_for_status()
                return response.json()

        except httpx.TimeoutException as e:
            logger.error(f"{provider_name} API timeout after {self.timeout}s")
            raise SearchError(
                f"LLM request timed out after {self.timeout} seconds. "
                "Try a simpler query or check your network connection."
            ) from e

        except httpx.HTTPStatusError as e:
            status_code = e.response.status_code
            error_msg = f"{provider_name} API error (HTTP {status_code})"

            # Try to get more details from the response
            try:
                error_body = e.response.json()
                error_detail = error_body.get("error", {}).get("message", "")
                if error_detail:
                    error_msg = f"{error_msg}: {error_detail}"
            except Exception:
                pass

            if status_code == 400:
                error_msg = f"{error_msg}. Check model name and request format."
            elif status_code == 401:
                env_var = (
                    "OPENAI_API_KEY"
                    if self.provider == "openai"
                    else "OPENROUTER_API_KEY"
                )
                error_msg = f"Invalid {provider_name} API key. Please check {env_var} environment variable."
            elif status_code == 429:
                error_msg = f"{provider_name} API rate limit exceeded. Please wait and try again."
            elif status_code >= 500:
                error_msg = f"{provider_name} API server error. Please try again later."

            logger.error(error_msg)
            raise SearchError(error_msg) from e

        except Exception as e:
            logger.error(f"{provider_name} API request failed: {e}")
            raise SearchError(f"LLM request failed: {e}") from e

    async def _bedrock_chat_completion(
        self, messages: list[dict[str, str]]
    ) -> dict[str, Any]:
        """Make chat completion request to AWS Bedrock using Converse API.

        Args:
            messages: List of message dictionaries with role and content

        Returns:
            API response dictionary in OpenAI format (for compatibility)

        Raises:
            SearchError: If Bedrock API request fails
        """
        try:
            # Convert messages to Bedrock format
            bedrock_messages = []
            for msg in messages:
                role = msg["role"]
                # Bedrock uses "user" and "assistant" (no "system" in messages)
                if role in ("user", "assistant"):
                    bedrock_messages.append(
                        {"role": role, "content": [{"text": msg["content"]}]}
                    )

            # Extract system message if present
            system_messages = [msg for msg in messages if msg["role"] == "system"]
            system_content = None
            if system_messages:
                system_content = [{"text": system_messages[0]["content"]}]

            # Build Bedrock request
            request_params = {
                "modelId": self.model,
                "messages": bedrock_messages,
                "inferenceConfig": {
                    "maxTokens": 4096,
                    "temperature": 0.7,
                },
            }

            # Add system message if present
            if system_content:
                request_params["system"] = system_content

            # Run boto3 call in executor (boto3 is synchronous)
            loop = asyncio.get_event_loop()
            bedrock_client = self._get_bedrock_client()

            response = await loop.run_in_executor(
                None,
                lambda: bedrock_client.converse(**request_params),
            )

            # Convert Bedrock response to OpenAI format for compatibility
            content = response["output"]["message"]["content"][0]["text"]

            return {
                "choices": [
                    {
                        "message": {
                            "role": "assistant",
                            "content": content,
                        },
                        "finish_reason": response.get("stopReason", "stop"),
                    }
                ],
                "usage": response.get("usage", {}),
            }

        except Exception as e:
            logger.error(f"Bedrock API request failed: {e}")
            error_msg = str(e)

            # Parse common Bedrock errors
            if "AccessDeniedException" in error_msg:
                error_msg = (
                    "AWS credentials invalid or insufficient permissions for Bedrock. "
                    "Check AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY."
                )
            elif "ValidationException" in error_msg:
                error_msg = f"Invalid Bedrock request: {error_msg}"
            elif "ThrottlingException" in error_msg:
                error_msg = (
                    "Bedrock API rate limit exceeded. Please wait and try again."
                )
            elif "ModelNotReadyException" in error_msg:
                error_msg = f"Bedrock model {self.model} is not ready or not available in your region."
            elif "ResourceNotFoundException" in error_msg:
                error_msg = (
                    f"Bedrock model {self.model} not found. Check model ID and region."
                )

            raise SearchError(f"Bedrock request failed: {error_msg}") from e

    def _format_results_for_analysis(self, search_results: dict[str, list[Any]]) -> str:
        """Format search results for LLM analysis.

        Args:
            search_results: Dictionary mapping search queries to their results

        Returns:
            Formatted string representation of results
        """
        formatted = []

        for i, (query, results) in enumerate(search_results.items(), 1):
            formatted.append(f"\n=== Query {i}: {query} ===")

            if not results:
                formatted.append("  No results found.")
                continue

            for j, result in enumerate(results[:5], 1):  # Top 5 per query
                # Extract key information from SearchResult
                file_path = str(result.file_path)
                similarity = result.similarity_score
                content_preview = result.content[:150].replace("\n", " ")

                formatted.append(
                    f"\n  Result {j}:\n"
                    f"    File: {file_path}\n"
                    f"    Similarity: {similarity:.3f}\n"
                    f"    Preview: {content_preview}..."
                )

                if result.function_name:
                    formatted.append(f"    Function: {result.function_name}")
                if result.class_name:
                    formatted.append(f"    Class: {result.class_name}")

        return "\n".join(formatted)

    def _parse_ranking_response(
        self,
        llm_response: str,
        search_results: dict[str, list[Any]],
        top_n: int,
    ) -> list[dict[str, Any]]:
        """Parse LLM ranking response into structured results.

        Args:
            llm_response: Raw LLM response text
            search_results: Original search results dictionary
            top_n: Maximum number of results to return

        Returns:
            List of ranked results with metadata
        """
        ranked = []
        current_result = {}

        for line in llm_response.split("\n"):
            line = line.strip()

            if line.startswith("RESULT:"):
                if current_result:
                    ranked.append(current_result)
                current_result = {"identifier": line.replace("RESULT:", "").strip()}

            elif line.startswith("RELEVANCE:"):
                current_result["relevance"] = line.replace("RELEVANCE:", "").strip()

            elif line.startswith("EXPLANATION:"):
                current_result["explanation"] = line.replace("EXPLANATION:", "").strip()

        # Add last result
        if current_result:
            ranked.append(current_result)

        # Map identifiers back to actual SearchResult objects
        enriched_results = []

        for item in ranked[:top_n]:
            identifier = item.get("identifier", "")

            # Parse identifier (e.g., "Query 1, Result 2" or "Query 1, Result 2 (filename.py)")
            try:
                parts = identifier.split(",")
                query_part = parts[0].replace("Query", "").strip()
                result_part = parts[1].replace("Result", "").strip()

                # Handle case where LLM includes filename in parentheses: "5 (config.py)"
                # Extract just the number
                query_match = re.match(r"(\d+)", query_part)
                result_match = re.match(r"(\d+)", result_part)

                if not query_match or not result_match:
                    logger.warning(
                        f"Could not extract numbers from identifier '{identifier}'"
                    )
                    continue

                query_idx = int(query_match.group(1)) - 1
                result_idx = int(result_match.group(1)) - 1

                # Get corresponding query and result
                queries = list(search_results.keys())
                if query_idx < len(queries):
                    query = queries[query_idx]
                    results = search_results[query]

                    if result_idx < len(results):
                        actual_result = results[result_idx]

                        enriched_results.append(
                            {
                                "result": actual_result,
                                "query": query,
                                "relevance": item.get("relevance", "Medium"),
                                "explanation": item.get(
                                    "explanation", "Relevant to query"
                                ),
                            }
                        )

            except (ValueError, IndexError) as e:
                logger.warning(f"Failed to parse result identifier '{identifier}': {e}")
                continue

        return enriched_results

    async def detect_intent(self, query: str) -> IntentType:
        """Detect user intent from query.

        Args:
            query: User's natural language query

        Returns:
            Intent type: "find", "answer", or "analyze"

        Raises:
            SearchError: If API call fails
        """
        system_prompt = """You are a code search intent classifier. Classify the user's query into ONE of these categories:

1. "find" - User wants to locate/search for something in the codebase
   Examples: "where is X", "find the function that", "show me the code for", "locate X"

2. "answer" - User wants an explanation/answer about the codebase
   Examples: "what does this do", "how does X work", "explain the architecture", "why is X used"

3. "analyze" - User wants analysis of code quality, metrics, complexity, or smells
   Examples: "what's complex", "code smells", "cognitive complexity", "quality issues",
   "dependencies", "coupling", "circular dependencies", "getting worse", "improving",
   "analyze the complexity", "find the worst code", "most complex functions"

Return ONLY the word "find", "answer", or "analyze" with no other text."""

        user_prompt = f"""Query: {query}

Intent:"""

        try:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ]

            response = await self._chat_completion(messages)

            content = (
                response.get("choices", [{}])[0].get("message", {}).get("content", "")
            )
            intent = content.strip().lower()

            if intent not in ("find", "answer", "analyze"):
                # Default to find if unclear
                logger.warning(
                    f"Unclear intent '{intent}' for query '{query}', defaulting to 'find'"
                )
                return "find"

            logger.debug(f"Detected intent '{intent}' for query: '{query}'")
            return intent  # type: ignore

        except Exception as e:
            logger.error(f"Failed to detect intent: {e}, defaulting to 'find'")
            return "find"

    async def stream_chat_completion(
        self, messages: list[dict[str, str]]
    ) -> AsyncIterator[str]:
        """Stream chat completion response chunk by chunk.

        Args:
            messages: List of message dictionaries with role and content

        Yields:
            Text chunks from the streaming response

        Raises:
            SearchError: If API request fails
        """
        # Route to Bedrock streaming if that's the provider
        if self.provider == "bedrock":
            async for chunk in self._bedrock_stream_chat_completion(messages):
                yield chunk
            return

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        if self.provider == "openrouter":
            headers["HTTP-Referer"] = "https://github.com/bobmatnyc/mcp-vector-search"
            headers["X-Title"] = "MCP Vector Search"

        payload = {
            "model": self.model,
            "messages": messages,
            "stream": True,
        }

        provider_name = self.provider.capitalize()

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                async with client.stream(
                    "POST", self.api_endpoint, headers=headers, json=payload
                ) as response:
                    response.raise_for_status()

                    async for line in response.aiter_lines():
                        line = line.strip()

                        # Skip empty lines and comments
                        if not line or line.startswith(":"):
                            continue

                        # Parse SSE format: "data: {json}"
                        if line.startswith("data: "):
                            data = line[6:]  # Remove "data: " prefix

                            # Check for end of stream
                            if data == "[DONE]":
                                break

                            try:
                                chunk = json.loads(data)
                                content = (
                                    chunk.get("choices", [{}])[0]
                                    .get("delta", {})
                                    .get("content")
                                )

                                if content:
                                    yield content

                            except json.JSONDecodeError as e:
                                logger.warning(f"Failed to parse SSE chunk: {e}")
                                continue

        except httpx.TimeoutException as e:
            logger.error(f"{provider_name} API timeout after {self.timeout}s")
            raise SearchError(
                f"LLM request timed out after {self.timeout} seconds. "
                "Try a simpler query or check your network connection."
            ) from e

        except httpx.HTTPStatusError as e:
            status_code = e.response.status_code
            error_msg = f"{provider_name} API error (HTTP {status_code})"

            if status_code == 401:
                env_var = (
                    "OPENAI_API_KEY"
                    if self.provider == "openai"
                    else "OPENROUTER_API_KEY"
                )
                error_msg = f"Invalid {provider_name} API key. Please check {env_var} environment variable."
            elif status_code == 429:
                error_msg = f"{provider_name} API rate limit exceeded. Please wait and try again."
            elif status_code >= 500:
                error_msg = f"{provider_name} API server error. Please try again later."

            logger.error(error_msg)
            raise SearchError(error_msg) from e

        except Exception as e:
            logger.error(f"{provider_name} streaming request failed: {e}")
            raise SearchError(f"LLM streaming failed: {e}") from e

    async def _bedrock_stream_chat_completion(
        self, messages: list[dict[str, str]]
    ) -> AsyncIterator[str]:
        """Stream chat completion from AWS Bedrock using Converse Stream API.

        Args:
            messages: List of message dictionaries with role and content

        Yields:
            Text chunks from the streaming response

        Raises:
            SearchError: If Bedrock streaming request fails
        """
        try:
            # Convert messages to Bedrock format
            bedrock_messages = []
            for msg in messages:
                role = msg["role"]
                if role in ("user", "assistant"):
                    bedrock_messages.append(
                        {"role": role, "content": [{"text": msg["content"]}]}
                    )

            # Extract system message if present
            system_messages = [msg for msg in messages if msg["role"] == "system"]
            system_content = None
            if system_messages:
                system_content = [{"text": system_messages[0]["content"]}]

            # Build Bedrock request
            request_params = {
                "modelId": self.model,
                "messages": bedrock_messages,
                "inferenceConfig": {
                    "maxTokens": 4096,
                    "temperature": 0.7,
                },
            }

            if system_content:
                request_params["system"] = system_content

            # Run streaming call in executor
            loop = asyncio.get_event_loop()
            bedrock_client = self._get_bedrock_client()

            response = await loop.run_in_executor(
                None,
                lambda: bedrock_client.converse_stream(**request_params),
            )

            # Process streaming response
            stream = response.get("stream")
            if stream:
                for event in stream:
                    if "contentBlockDelta" in event:
                        delta = event["contentBlockDelta"]["delta"]
                        if "text" in delta:
                            yield delta["text"]

        except Exception as e:
            logger.error(f"Bedrock streaming request failed: {e}")
            error_msg = str(e)

            if "AccessDeniedException" in error_msg:
                error_msg = (
                    "AWS credentials invalid or insufficient permissions for Bedrock."
                )
            elif "ThrottlingException" in error_msg:
                error_msg = (
                    "Bedrock API rate limit exceeded. Please wait and try again."
                )

            raise SearchError(f"Bedrock streaming failed: {error_msg}") from e

    async def generate_answer(
        self,
        query: str,
        context: str,
        conversation_history: list[dict[str, str]] | None = None,
    ) -> str:
        """Generate answer to user question using codebase context.

        Args:
            query: User's question
            context: Relevant code context from search results
            conversation_history: Previous conversation messages (optional)

        Returns:
            LLM response text

        Raises:
            SearchError: If API call fails
        """
        system_prompt = f"""You are a helpful code assistant analyzing a codebase. Answer the user's questions based on the provided code context.

Code Context:
{context}

Guidelines:
- Be concise but thorough in explanations
- Reference specific functions, classes, or files when relevant
- Use code examples from the context when helpful
- If the context doesn't contain enough information, say so
- Use markdown formatting for code snippets"""

        messages = [{"role": "system", "content": system_prompt}]

        # Add conversation history if provided
        if conversation_history:
            messages.extend(conversation_history)

        # Add current query
        messages.append({"role": "user", "content": query})

        try:
            response = await self._chat_completion(messages)
            content = (
                response.get("choices", [{}])[0].get("message", {}).get("content", "")
            )

            logger.debug(f"Generated answer for query: '{query}'")
            return content

        except Exception as e:
            logger.error(f"Failed to generate answer: {e}")
            raise SearchError(f"Failed to generate answer: {e}") from e

    async def chat_with_tools(
        self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]
    ) -> dict[str, Any]:
        """Chat completion with tool/function calling support.

        Args:
            messages: List of message dictionaries
            tools: List of tool definitions

        Returns:
            API response with tool calls or final message

        Raises:
            SearchError: If API request fails

        Note:
            Bedrock tool calling is not yet implemented. Falls back to regular chat.
        """
        # TODO: Implement Bedrock tool calling when needed
        # Bedrock uses a different tool format (toolConfig) than OpenAI
        if self.provider == "bedrock":
            logger.warning(
                "Tool calling not yet implemented for Bedrock, falling back to regular chat"
            )
            return await self._chat_completion(messages)

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        if self.provider == "openrouter":
            headers["HTTP-Referer"] = "https://github.com/bobmatnyc/mcp-vector-search"
            headers["X-Title"] = "MCP Vector Search"

        payload = {
            "model": self.model,
            "messages": messages,
            "tools": tools,
            "tool_choice": "auto",
        }

        provider_name = self.provider.capitalize()

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    self.api_endpoint,
                    headers=headers,
                    json=payload,
                )

                response.raise_for_status()
                return response.json()

        except httpx.TimeoutException as e:
            logger.error(f"{provider_name} API timeout after {self.timeout}s")
            raise SearchError(
                f"LLM request timed out after {self.timeout} seconds."
            ) from e

        except httpx.HTTPStatusError as e:
            status_code = e.response.status_code
            error_msg = f"{provider_name} API error (HTTP {status_code})"

            # Try to get more details from the response
            try:
                error_body = e.response.json()
                error_detail = error_body.get("error", {}).get("message", "")
                if error_detail:
                    error_msg = f"{error_msg}: {error_detail}"
            except Exception:
                pass

            if status_code == 400:
                error_msg = f"{error_msg}. Check model name and request format."
            elif status_code == 401:
                env_var = (
                    "OPENAI_API_KEY"
                    if self.provider == "openai"
                    else "OPENROUTER_API_KEY"
                )
                error_msg = f"Invalid {provider_name} API key. Check {env_var}."
            elif status_code == 429:
                error_msg = f"{provider_name} API rate limit exceeded."
            elif status_code >= 500:
                error_msg = f"{provider_name} API server error."

            logger.error(error_msg)
            raise SearchError(error_msg) from e

        except Exception as e:
            logger.error(f"{provider_name} API request failed: {e}")
            raise SearchError(f"LLM request failed: {e}") from e
