"""
Core SQL Agent — orchestrates schema retrieval, GPT-4 generation,
dry run validation, caching, and execution.
"""

import json
import hashlib
from openai import OpenAI
from datanaut.schema.retriever import SchemaRetriever
from datanaut.utils.bigquery_client import BigQueryClient
from datanaut.utils.cache import QueryCache
from datanaut.utils.display import print_step


class SQLAgent:
    def __init__(self):
        self.client = OpenAI()  # reads OPENAI_API_KEY from env
        self.schema_retriever = SchemaRetriever()
        self.bq = BigQueryClient()
        self.cache = QueryCache()
        self.max_retries = 3

    # ------------------------------------------------------------------ #
    #  Main entry point                                                    #
    # ------------------------------------------------------------------ #

    def run(self, question: str, conversation_history: list) -> dict:
        """
        Orchestrates the full pipeline for a single user question.
        Returns a result dict with keys: success, sql, data, row_count,
        bytes_estimate, from_cache, error.
        """

        # Step 1: retrieve relevant schema
        print_step("Retrieving relevant schema...")
        schema_context = self.schema_retriever.get_relevant_schema(question)

        # Step 2: generate SQL with GPT-4
        print_step("Generating SQL with GPT-4...")
        sql, generation_error = self._generate_sql(
            question, schema_context, conversation_history
        )
        if generation_error:
            return {"success": False, "error": generation_error}

        # Step 3: validate + get cost estimate via dry run
        print_step("Running BigQuery dry run...")
        bytes_estimate, dry_run_error = self.bq.dry_run(sql)
        if dry_run_error:
            # Retry loop — pass the error back to GPT-4 for self-correction
            for attempt in range(self.max_retries):
                print_step(f"Retrying SQL generation (attempt {attempt + 1}/{self.max_retries})...")
                sql, generation_error = self._generate_sql(
                    question, schema_context, conversation_history,
                    previous_sql=sql, error=dry_run_error
                )
                if generation_error:
                    return {"success": False, "error": generation_error}

                bytes_estimate, dry_run_error = self.bq.dry_run(sql)
                if not dry_run_error:
                    break
            else:
                return {
                    "success": False,
                    "error": f"Could not generate valid SQL after {self.max_retries} attempts.\nLast error: {dry_run_error}"
                }

        # Step 4: check cache
        cache_key = hashlib.md5(sql.strip().lower().encode()).hexdigest()
        cached = self.cache.get(cache_key)
        if cached:
            return {
                "success": True,
                "sql": sql,
                "bytes_estimate": self._format_bytes(bytes_estimate),
                "data": cached["data"],
                "row_count": cached["row_count"],
                "from_cache": True,
            }

        # Step 5: execute (caller confirms before this is reached)
        print_step("Executing query on BigQuery...")
        data, row_count, exec_error = self.bq.execute(sql)
        if exec_error:
            return {"success": False, "error": exec_error}

        # Cache successful result
        self.cache.set(cache_key, {"data": data, "row_count": row_count})

        return {
            "success": True,
            "sql": sql,
            "bytes_estimate": self._format_bytes(bytes_estimate),
            "data": data,
            "row_count": row_count,
            "from_cache": False,
        }

    # ------------------------------------------------------------------ #
    #  GPT-4 SQL generation                                               #
    # ------------------------------------------------------------------ #

    def _generate_sql(
        self,
        question: str,
        schema_context: str,
        conversation_history: list,
        previous_sql: str = None,
        error: str = None,
    ) -> tuple[str, str]:
        """
        Calls GPT-4 with function calling to generate optimized BigQuery SQL.
        Returns (sql_string, error_string). One of the two will be None.
        """

        system_prompt = self._build_system_prompt(schema_context)
        messages = self._build_messages(
            question, conversation_history, previous_sql, error
        )

        tools = [
            {
                "type": "function",
                "function": {
                    "name": "generate_bigquery_sql",
                    "description": "Generate an optimized BigQuery SQL query for the user's question.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "sql": {
                                "type": "string",
                                "description": "The complete BigQuery SQL query. Must be valid Standard SQL."
                            },
                            "explanation": {
                                "type": "string",
                                "description": "One-sentence explanation of what the query does."
                            },
                            "optimizations_applied": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "List of BigQuery optimizations applied (e.g. partition pruning, approx functions)."
                            }
                        },
                        "required": ["sql", "explanation"]
                    }
                }
            }
        ]

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o",
                messages=[{"role": "system", "content": system_prompt}] + messages,
                tools=tools,
                tool_choice={"type": "function", "function": {"name": "generate_bigquery_sql"}},
                temperature=0.1,  # low temp for deterministic SQL
            )

            tool_call = response.choices[0].message.tool_calls[0]
            args = json.loads(tool_call.function.arguments)
            sql = args.get("sql", "").strip()

            if not sql:
                return None, "GPT-4 returned an empty SQL query."

            # Print explanation and optimizations as context for the user
            explanation = args.get("explanation", "")
            optimizations = args.get("optimizations_applied", [])
            if explanation:
                print(f"  GPT-4: {explanation}")
            if optimizations:
                print(f"  Optimizations: {', '.join(optimizations)}")

            return sql, None

        except Exception as e:
            return None, f"GPT-4 error: {str(e)}"

    def _build_system_prompt(self, schema_context: str) -> str:
        return f"""You are an expert BigQuery SQL generator. Your job is to convert natural language questions into optimized BigQuery Standard SQL queries.

BIGQUERY SCHEMA:
{schema_context}

STRICT RULES — follow all of these:
1. ALWAYS filter on partition columns (e.g. WHERE date >= DATE_SUB(CURRENT_DATE(), INTERVAL 30 DAY))
2. NEVER use SELECT * — always list specific columns needed
3. Use APPROX_COUNT_DISTINCT instead of COUNT(DISTINCT ...) for large datasets
4. Use DATE_TRUNC or TIMESTAMP_TRUNC for date grouping
5. Prefer CTEs (WITH clauses) over subqueries for readability
6. Always qualify table names with dataset: `dataset.table_name`
7. Use LIMIT for exploratory queries (default 1000 rows unless user specifies)
8. For aggregations, always include relevant GROUP BY columns
9. Use SAFE_DIVIDE instead of / to avoid division by zero
10. Add ORDER BY for sorted results

When correcting a previous SQL error, fix ONLY the problematic part — preserve working logic.
"""

    def _build_messages(
        self,
        question: str,
        conversation_history: list,
        previous_sql: str = None,
        error: str = None,
    ) -> list:
        messages = []

        # Inject last 3 turns of conversation for context
        for turn in conversation_history[-3:]:
            messages.append({"role": "user", "content": turn["question"]})
            messages.append({
                "role": "assistant",
                "content": f"Generated SQL:\n```sql\n{turn['sql']}\n```"
            })

        # Current question
        if previous_sql and error:
            messages.append({
                "role": "user",
                "content": (
                    f"Question: {question}\n\n"
                    f"Previous SQL that failed:\n```sql\n{previous_sql}\n```\n\n"
                    f"BigQuery error: {error}\n\n"
                    f"Please fix the SQL."
                )
            })
        else:
            messages.append({"role": "user", "content": f"Question: {question}"})

        return messages

    # ------------------------------------------------------------------ #
    #  Helpers                                                             #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _format_bytes(num_bytes: int) -> str:
        if num_bytes is None:
            return "unknown"
        for unit in ["B", "KB", "MB", "GB", "TB"]:
            if num_bytes < 1024:
                return f"{num_bytes:.1f} {unit}"
            num_bytes /= 1024
        return f"{num_bytes:.1f} PB"
