"""Query refiner using LLM for intelligent query correction."""

import json
import re
from typing import Any, Union

from openai import AsyncOpenAI, AsyncAzureOpenAI

from .prompts import (
    SYSTEM_PROMPT,
    INITIAL_QUERY_PROMPT,
    REFINEMENT_PROMPT,
    SCHEMA_ANALYSIS_PROMPT,
    EXPLANATION_ANALYSIS_PROMPT,
    RESPONSE_ANALYSIS_PROMPT,
    EXPLANATION_VERIFICATION_PROMPT,
    RESPONSE_VERIFICATION_PROMPT,
    QUERY_TEMPLATE_VERIFICATION_PROMPT,
    EXPLANATION_PATH_SELECTION_PROMPT,
)
from .models import SchemaInfo


class QueryRefiner:
    """Uses LLM to generate and refine BTQL queries."""

    def __init__(
        self,
        client: Union[AsyncOpenAI, AsyncAzureOpenAI],
        model: str,
    ) -> None:
        """Initialize QueryRefiner with an OpenAI-compatible client.

        Args:
            client: Any OpenAI-compatible async client (AsyncOpenAI or AsyncAzureOpenAI)
            model: Model name or deployment to use
        """
        self.client = client
        self.model = model

    async def generate_initial_query(
        self,
        experiment_id: str,
        metrics: list[str],
        filters: list[str] | None = None,
    ) -> str:
        """Generate an initial BTQL query using LLM.

        Args:
            experiment_id: The experiment ID to query
            metrics: List of metrics to fetch
            filters: Optional list of filter conditions

        Returns:
            Generated BTQL query string
        """
        filters_str = " and ".join(filters) if filters else "None"
        metrics_str = ", ".join(metrics)

        user_prompt = INITIAL_QUERY_PROMPT.format(
            experiment_id=experiment_id,
            metrics=metrics_str,
            filters=filters_str,
        )

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )

        query = response.choices[0].message.content.strip()
        # Remove markdown code blocks if present
        query = self._clean_query(query)
        return query

    async def refine_query(
        self,
        experiment_id: str,
        metrics: list[str],
        previous_query: str,
        error_message: str,
        attempt_history: list[dict] | None = None,
        filters: list[str] | None = None,
    ) -> str:
        """Refine a failed BTQL query based on the error.

        Args:
            experiment_id: The experiment ID
            metrics: List of metrics to fetch
            previous_query: The query that failed
            error_message: The error message from the failed query
            attempt_history: List of previous attempts with their errors
            filters: Optional list of filter conditions

        Returns:
            Refined BTQL query string
        """
        filters_str = " and ".join(filters) if filters else "None"
        metrics_str = ", ".join(metrics)

        # Format attempt history
        history_str = "None"
        if attempt_history:
            history_parts = []
            for i, attempt in enumerate(attempt_history, 1):
                history_parts.append(
                    f"Attempt {i}:\n  Query: {attempt['query']}\n  Error: {attempt['error']}"
                )
            history_str = "\n".join(history_parts)

        user_prompt = REFINEMENT_PROMPT.format(
            experiment_id=experiment_id,
            metrics=metrics_str,
            filters=filters_str,
            previous_query=previous_query,
            error_message=error_message,
            attempt_history=history_str,
        )

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )

        query = response.choices[0].message.content.strip()
        query = self._clean_query(query)
        return query

    def _clean_query(self, query: str) -> str:
        """Remove markdown formatting from query."""
        # Remove code blocks
        if query.startswith("```"):
            lines = query.split("\n")
            # Remove first line (```sql or ```)
            lines = lines[1:]
            # Remove last line if it's ```
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            query = "\n".join(lines)

        return query.strip()

    async def analyze_schema(
        self,
        sample_data: list[dict[str, Any]],
        attempt_history: list[dict[str, str]] | None = None,
        provided_metrics: list[str] | None = None,
    ) -> SchemaInfo:
        """Analyze sample data to discover schema structure.

        Args:
            sample_data: 1-2 sample rows from the experiment
            attempt_history: Optional list of previous attempts with their errors/issues
            provided_metrics: Optional list of metrics already known (skips metric discovery)

        Returns:
            SchemaInfo with discovered paths to scores and explanations
        """
        # Truncate large fields to avoid token limits
        truncated_data = self._truncate_sample(sample_data)
        sample_json = json.dumps(truncated_data, indent=4, default=str)

        # Format attempt history section
        attempt_history_section = ""
        if attempt_history:
            attempt_history_section = "PREVIOUS FAILED ATTEMPTS:\n"
            attempt_history_section += "The following attempts produced incorrect paths. Learn from these errors and find DIFFERENT paths:\n\n"
            for i, attempt in enumerate(attempt_history, 1):
                attempt_history_section += f"Attempt {i}:\n"
                attempt_history_section += f"  Query template used: {attempt.get('query', 'N/A')}\n"
                attempt_history_section += f"  Issue: {attempt.get('error', 'N/A')}\n\n"
            attempt_history_section += "IMPORTANT: Do NOT use the same paths as the failed attempts. Search for alternative locations in the data structure.\n"

        # Format metrics section - if provided, skip metric discovery
        if provided_metrics:
            metrics_section = f"METRICS (already provided, do NOT discover from scores):\n{', '.join(provided_metrics)}"
            task_step_1 = ""  # Skip step 1
        else:
            metrics_section = ""
            task_step_1 = "- Identify available metric names in the scores object\n"

        user_prompt = SCHEMA_ANALYSIS_PROMPT.format(
            sample_data=sample_json,
            attempt_history_section=attempt_history_section,
            metrics_section=metrics_section,
            task_step_1=task_step_1,
        )

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data schema analyst. Analyze JSON data structures and identify patterns."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=2000,
        )

        content = response.choices[0].message.content.strip()
        schema_dict = self._parse_schema_response(content)

        return SchemaInfo(
            scores_path=schema_dict.get("scores_path", "scores"),
            available_metrics=schema_dict.get("available_metrics", []),
            explanation_paths=schema_dict.get("explanation_paths", []),
            id_field=schema_dict.get("id_field", "root_span_id"),
            raw_sample=sample_data[0] if sample_data else None,
        )

    async def analyze_explanations(
        self,
        sample_data: list[dict[str, Any]],
        metrics: list[str],
        attempt_history: list[dict[str, str]] | None = None,
    ) -> dict[str, Any]:
        """Analyze sample data to discover paths to metric explanations.

        Args:
            sample_data: 1-2 sample rows from the experiment
            metrics: List of metrics to find explanations for
            attempt_history: Optional list of previous attempts with their errors/issues

        Returns:
            Dictionary with scores_path, available_metrics, explanation_paths, id_field
        """
        # Truncate large fields to avoid token limits
        truncated_data = self._truncate_sample(sample_data)
        sample_json = json.dumps(truncated_data, indent=4, default=str)

        # Format attempt history section
        attempt_history_section = ""
        if attempt_history:
            attempt_history_section = "PREVIOUS FAILED ATTEMPTS:\n"
            attempt_history_section += "The following attempts produced incorrect paths. Learn from these errors and find DIFFERENT paths:\n\n"
            for i, attempt in enumerate(attempt_history, 1):
                attempt_history_section += f"Attempt {i}:\n"
                attempt_history_section += f"  Query template used: {attempt.get('query', 'N/A')}\n"
                attempt_history_section += f"  Issue: {attempt.get('error', 'N/A')}\n\n"
            attempt_history_section += "IMPORTANT: Do NOT use the same paths as the failed attempts. Search for alternative locations in the data structure.\n"

        metrics_str = ", ".join(metrics)
        
        user_prompt = EXPLANATION_ANALYSIS_PROMPT.format(
            sample_data=sample_json,
            attempt_history_section=attempt_history_section,
            metrics=metrics_str,
        )

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data schema analyst. Analyze JSON data structures and identify patterns for metric explanations."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=2000,
        )

        content = response.choices[0].message.content.strip()
        return self._parse_schema_response(content)

    async def analyze_response_path(
        self,
        sample_data: list[dict[str, Any]],
        attempt_history: list[dict[str, str]] | None = None,
    ) -> dict[str, Any]:
        """Analyze sample data to discover path to AI model's response.

        Args:
            sample_data: 1-2 sample rows from the experiment
            attempt_history: Optional list of previous attempts with their errors/issues

        Returns:
            Dictionary with response_field_path
        """
        # Truncate large fields to avoid token limits
        truncated_data = self._truncate_sample(sample_data)
        sample_json = json.dumps(truncated_data, indent=4, default=str)

        # Format attempt history section
        attempt_history_section = ""
        if attempt_history:
            attempt_history_section = "PREVIOUS FAILED ATTEMPTS:\n"
            attempt_history_section += "The following attempts produced incorrect paths. Learn from these errors and find DIFFERENT paths:\n\n"
            for i, attempt in enumerate(attempt_history, 1):
                attempt_history_section += f"Attempt {i}:\n"
                attempt_history_section += f"  Query template used: {attempt.get('query', 'N/A')}\n"
                attempt_history_section += f"  Issue: {attempt.get('error', 'N/A')}\n\n"
            attempt_history_section += "IMPORTANT: Do NOT use the same paths as the failed attempts. Search for alternative locations in the data structure.\n"

        user_prompt = RESPONSE_ANALYSIS_PROMPT.format(
            sample_data=sample_json,
            attempt_history_section=attempt_history_section,
        )

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data schema analyst. Analyze JSON data structures and identify patterns for AI model responses."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )

        content = response.choices[0].message.content.strip()
        return self._parse_schema_response(content)

    async def verify_explanation_query(
        self,
        sample_result: dict[str, Any],
        available_metrics: list[str],
    ) -> dict[str, Any]:
        """Verify a query result row has correct explanation content using LLM analysis.
        
        Args:
            sample_result: Single row from executing the explanation query
            available_metrics: List of metrics the query should support
            
        Returns:
            Dictionary with verification results:
            - is_valid: bool
            - explanation: explanation of why valid or what's wrong
        """
        # Truncate sample result to avoid token limits
        # truncated_result = self._truncate_sample([sample_result], max_str_len=2000)
        sample_json = json.dumps(sample_result, indent=2, default=str)
        
        user_prompt = EXPLANATION_VERIFICATION_PROMPT.format(
            sample_result=sample_json,
            available_metrics=", ".join(available_metrics),
        )
        
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data quality analyst. Verify query results for explanation extraction."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )
        
        content = response.choices[0].message.content.strip()
        result = self._parse_schema_response(content)
        
        # Ensure required fields exist with defaults
        return {
            "is_valid": result.get("is_valid", False),
            "explanation": result.get("explanation", "No explanation provided")
        }

    async def verify_response_query(
        self,
        sample_result: dict[str, Any],
    ) -> dict[str, Any]:
        """Verify a query result row has correct response content using LLM analysis.
        
        Args:
            sample_result: Single row from executing the response query
            
        Returns:
            Dictionary with verification results:
            - is_valid: bool
            - explanation: explanation of why valid or what's wrong
        """
        # Truncate sample result to avoid token limits
        truncated_result = self._truncate_sample([sample_result], max_str_len=2000)
        sample_json = json.dumps(truncated_result[0] if truncated_result else {}, indent=2, default=str)
        
        user_prompt = RESPONSE_VERIFICATION_PROMPT.format(
            sample_result=sample_json,
        )
        
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data quality analyst. Verify query results for AI response extraction."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )
        
        content = response.choices[0].message.content.strip()
        result = self._parse_schema_response(content)
        
        # Ensure required fields exist with defaults
        return {
            "is_valid": result.get("is_valid", False),
            "explanation": result.get("explanation", "No explanation provided")
        }

    def _truncate_sample(self, data: list[dict[str, Any]], max_str_len: int = 3000) -> list[dict[str, Any]]:
        """Truncate long string values in sample data to reduce tokens."""
        def truncate_value(v: Any, key: str = "") -> Any:
            # Special handling for fields to ignore - replace with placeholder
            if key == "inferences":
                return "...[inferences truncated]"
            if key == "search_results":
                return "...[search_results truncated]"
            
            if isinstance(v, str) and len(v) > max_str_len:
                return v[:max_str_len] + "...[truncated]"
            elif isinstance(v, dict):
                return {k: truncate_value(val, k) for k, val in v.items()}
            elif isinstance(v, list):
                # Only keep first 2 items for other arrays
                return [truncate_value(item, key) for item in v[:2]]
            return v

        return [truncate_value(row) for row in data[:4]]  # Keep up to 4 rows

    def _parse_schema_response(self, content: str) -> dict[str, Any]:
        """Parse JSON from LLM response."""
        # Try direct JSON parse
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        # Try to extract JSON from content
        json_match = re.search(r'\{[\s\S]*\}', content)
        if json_match:
            try:
                return json.loads(json_match.group(0))
            except json.JSONDecodeError:
                pass

        # Return empty defaults on failure
        return {}

    def _convert_path_to_btql(self, path: str) -> str:
        """Convert dot-notation path to BTQL format with brackets for array indices.
        
        e.g., 'input.output.0.picasso_ai_response' -> 'input.output[0].picasso_ai_response'
        """
        parts = path.split(".")
        result_parts = []
        i = 0
        while i < len(parts):
            part = parts[i]
            if part.isdigit():
                # This is an array index, append to previous part with brackets
                if result_parts:
                    result_parts[-1] = f"{result_parts[-1]}[{part}]"
            else:
                result_parts.append(part)
            i += 1
        return ".".join(result_parts)

    async def verify_query_template(
        self,
        sample_result: dict[str, Any],
        available_metrics: list[str],
    ) -> dict[str, Any]:
        """Verify a query result row has correct content using LLM analysis.
        
        Args:
            sample_result: Single row from executing the query
            available_metrics: List of metrics the query should support
            
        Returns:
            Dictionary with verification results:
            - is_valid: bool
            - explanation: explanation of why valid or what's wrong
        """
        # Truncate sample result to avoid token limits
        truncated_result = self._truncate_sample([sample_result], max_str_len=2000)
        sample_json = json.dumps(truncated_result[0] if truncated_result else {}, indent=2, default=str)
        
        user_prompt = QUERY_TEMPLATE_VERIFICATION_PROMPT.format(
            sample_result=sample_json,
            available_metrics=", ".join(available_metrics),
        )
        
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data quality analyst. Verify query results match expected schema."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )
        
        content = response.choices[0].message.content.strip()
        result = self._parse_schema_response(content)
        
        # Ensure required fields exist with defaults
        return {
            "is_valid": result.get("is_valid", False),
            "explanation": result.get("explanation", "No explanation provided")
        }

    async def select_best_explanation_paths(
        self,
        candidates: list[dict[str, Any]],
        available_metrics: list[str],
    ) -> dict[str, Any]:
        """Select the best explanation path configuration from multiple candidates.
        
        This is a fallback when all attempts fail verification but some executed successfully.
        The LLM selects the candidate with best metric coverage and quality.
        
        Args:
            candidates: List of candidates, each containing:
                - index: The attempt index
                - explanation_paths: The discovered paths
                - query_result: Sample data from executing the query
                - query_template: The generated query template
            available_metrics: List of metrics that should be covered
            
        Returns:
            Dictionary with:
            - selected_index: Index of the best candidate
            - explanation: Why this candidate was selected
            - metrics_coverage: Coverage information
        """
        if not candidates:
            return {
                "selected_index": -1,
                "explanation": "No candidates provided",
                "metrics_coverage": {"selected_candidate_metrics_found": [], "coverage_percentage": 0}
            }
        
        if len(candidates) == 1:
            # Only one candidate, return it directly
            return {
                "selected_index": 0,
                "explanation": "Only one candidate available",
                "metrics_coverage": {"selected_candidate_metrics_found": [], "coverage_percentage": 0}
            }
        
        # Format candidates info for the prompt
        candidates_info_parts = []
        for i, candidate in enumerate(candidates):
            # Truncate query result to avoid token limits
            truncated_result = self._truncate_sample(
                candidate.get("query_result", [])[:2], 
                max_str_len=1500
            )
            sample_json = json.dumps(truncated_result, indent=2, default=str)
            
            candidates_info_parts.append(
                f"=== Candidate {i} ===\n"
                f"Explanation paths: {json.dumps(candidate.get('explanation_paths', []))}\n"
                f"Sample extracted data:\n{sample_json}\n"
            )
        
        candidates_info = "\n".join(candidates_info_parts)
        
        user_prompt = EXPLANATION_PATH_SELECTION_PROMPT.format(
            available_metrics=", ".join(available_metrics),
            candidates_info=candidates_info,
        )
        
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data quality analyst. Select the best configuration based on coverage and quality."},
                {"role": "user", "content": user_prompt},
            ],
            max_completion_tokens=1000,
        )
        
        content = response.choices[0].message.content.strip()
        result = self._parse_schema_response(content)
        
        # Ensure selected_index is valid
        selected_index = result.get("selected_index", 0)
        if not isinstance(selected_index, int) or selected_index < 0 or selected_index >= len(candidates):
            selected_index = 0
        
        return {
            "selected_index": selected_index,
            "explanation": result.get("explanation", "No explanation provided"),
            "metrics_coverage": result.get("metrics_coverage", {
                "selected_candidate_metrics_found": [],
                "coverage_percentage": 0
            })
        }
