"""
Annie SDK Agent

The main entry point for the Annie SDK. The Agent handles communication
with the Annie API and local SQL execution via connectors.

Flow:
1. User calls agent.run("natural language query")
2. Agent sends NL query + schema context to Annie API
3. Annie API returns DSL + SQL in the connector's dialect
4. Agent executes SQL locally via the connector
5. Results are formatted as AgentResponse
"""

from __future__ import annotations

import logging
import os
import random
from decimal import Decimal
from typing import Any

import httpx

from .connectors.base import BaseConnector
from .exceptions import APIError, AuthenticationError, RateLimitError, TimeoutError, ValidationError
from .response import AgentResponse, ChartFormat

logger = logging.getLogger(__name__)

DEFAULT_API_URL = "https://chat.pandas-ai.com"
DEFAULT_TIMEOUT = 60

DEFAULT_DENIAL_MESSAGES = [
    "I can only help with questions about your data. Try asking something like 'Show revenue by month'.",
    "That doesn't seem to be a data question. Try asking about your data, like 'Top 10 customers by sales'.",
    "I'm a data assistant. Ask me about your data — for example, 'What's the average order value?'",
]


class Agent:
    """Annie SDK Agent.

    Handles natural language queries against your database using the
    Annie API for NL-to-SQL conversion and local execution for data privacy.

    Args:
        connector: Database connector instance
        api_key: Annie API key (or set ANNIE_API_KEY env var)
        api_url: Annie API URL (or set ANNIE_API_URL env var)
        denial_messages: Custom messages for non-data queries

    Example:
        ```python
        from sdk import Agent, PostgresConnector, ConnectorTable

        connector = PostgresConnector(
            connection_string="postgresql://user:pass@localhost/db",
            tables=[ConnectorTable(name="orders", description="Customer orders")],
        )

        with Agent(connector=connector, api_key="your-key") as agent:
            response = agent.run("Show me top 10 orders by amount")
            print(response.data)
        ```
    """

    def __init__(
        self,
        connector: BaseConnector,
        api_key: str | None = None,
        api_url: str | None = None,
        denial_messages: list[str] | None = None,
    ):
        # Resolve API key
        self._api_key = api_key or os.environ.get("ANNIE_API_KEY")
        if not self._api_key:
            raise AuthenticationError(
                "API key is required. Pass api_key parameter or set ANNIE_API_KEY environment variable."
            )

        # Resolve API URL
        self._api_url = (
            api_url or os.environ.get("ANNIE_API_URL") or DEFAULT_API_URL
        ).rstrip("/")

        self._connector = connector
        self._denial_messages = denial_messages or DEFAULT_DENIAL_MESSAGES

        # Validate connector has tables
        if not connector.tables:
            raise ValidationError(
                "Connector must have at least one table configured",
                field="connector.tables",
            )

        # Initialize HTTP client
        self._client = httpx.Client(
            base_url=self._api_url,
            headers={
                "Authorization": f"Bearer {self._api_key}",
                "Content-Type": "application/json",
            },
            timeout=DEFAULT_TIMEOUT,
        )

    def run(self, query: str, *, explain: bool = False) -> AgentResponse:
        """Run a natural language query against your database.

        Args:
            query: Natural language query (e.g., "Show me top 10 customers by revenue")
            explain: If True, also generate AI insights about the results

        Returns:
            AgentResponse with data, SQL, DSL, and optional insights
        """
        if not query or not query.strip():
            return AgentResponse.from_error("Query cannot be empty")

        # 1. Build schema context from connector
        schema_context = self._connector.get_schema_context()

        # 2. Call Annie API: NL → DSL + SQL
        api_response = self._call_query_api(query, schema_context)

        # 3. Handle non-data responses (guardrail denials)
        if api_response.get("type") == "text":
            message = api_response.get("message") or random.choice(self._denial_messages)
            response = AgentResponse(
                success=True,
                visualization="table",
            )
            response.add_text(message)
            return response

        # 4. Handle API errors
        if not api_response.get("success", True):
            error_msg = api_response.get("message", "API returned an error")
            return AgentResponse.from_error(error_msg)

        sql = api_response.get("sql")
        dsl = api_response.get("dsl")
        visualization = api_response.get("visualization", "table")

        if not sql:
            return AgentResponse.from_error("API did not return SQL")

        # 5. Execute SQL locally via connector
        try:
            if not self._connector.is_connected():
                self._connector.connect()

            results = self._connector.execute(sql)
        except Exception as e:
            response = AgentResponse(
                sql=sql,
                dsl=dsl,
                visualization=visualization,
                success=False,
            )
            response.add_error(f"Query execution failed: {e}")
            return response

        # 6. Format results as AgentResponse
        response = AgentResponse(
            sql=sql,
            dsl=dsl,
            visualization=visualization,
            success=True,
        )

        if results:
            columns = list(results[0].keys())
            rows = [list(row.values()) for row in results]
            response.add_data(columns=columns, rows=rows)
        else:
            response.add_data(columns=[], rows=[])

        # 6b. Generate chart from results
        if results and visualization not in ("table", "kpi"):
            chart_config = self._generate_chart(results, dsl or {}, visualization)
            if chart_config:
                response.add_chart(
                    content=chart_config["content"],
                    format=ChartFormat(chart_config.get("format", "chartjs")),
                    title=chart_config.get("title"),
                )

        # 7. Optionally get AI explanation
        if explain and results:
            try:
                explanation = self._call_explain_api(
                    query=query,
                    results=results[:50],
                    sql=sql,
                    dsl=dsl,
                    schema=schema_context,
                )
                if explanation:
                    response.add_text(explanation)
            except Exception as e:
                logger.warning(f"Failed to get explanation: {e}")
                # Don't fail the whole response if explain fails

        return response

    def close(self) -> None:
        """Close the agent and its connector."""
        try:
            self._connector.disconnect()
        except Exception:
            pass
        try:
            self._client.close()
        except Exception:
            pass

    # =========================================================================
    # Context manager
    # =========================================================================

    def __enter__(self) -> Agent:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    # =========================================================================
    # Private methods
    # =========================================================================

    def _call_query_api(
        self,
        query: str,
        schema: dict[str, Any],
    ) -> dict[str, Any]:
        """Call the Annie /sdk/query endpoint."""
        payload = {
            "query": query,
            "schema": schema,
            "dialect": self._connector.dialect,
        }

        try:
            response = self._client.post("/sdk/query", json=payload)
        except httpx.TimeoutException as e:
            raise TimeoutError(f"API request timed out: {e}")
        except httpx.ConnectError as e:
            raise APIError(f"Failed to connect to Annie API at {self._api_url}: {e}")

        if response.status_code == 401:
            raise AuthenticationError("Invalid API key")
        if response.status_code == 429:
            retry_after = response.headers.get("retry-after")
            raise RateLimitError(
                "Rate limit exceeded",
                retry_after=int(retry_after) if retry_after else None,
            )
        if response.status_code >= 500:
            raise APIError(
                f"Annie API server error: {response.status_code}",
                status_code=response.status_code,
            )

        try:
            return response.json()
        except Exception as e:
            raise APIError(f"Failed to parse API response: {e}")

    def _call_explain_api(
        self,
        query: str,
        results: list[dict[str, Any]],
        sql: str | None,
        dsl: dict[str, Any] | None,
        schema: dict[str, Any] | None,
    ) -> str | None:
        """Call the Annie /sdk/explain endpoint."""
        payload = {
            "query": query,
            "results": results,
            "sql": sql,
            "dsl": dsl,
            "schema": schema,
        }

        try:
            response = self._client.post("/sdk/explain", json=payload)

            if response.status_code != 200:
                logger.warning(f"Explain API returned {response.status_code}")
                return None

            data = response.json()
            return data.get("explanation") if data.get("success") else None

        except Exception as e:
            logger.warning(f"Explain API call failed: {e}")
            return None

    def _generate_chart(
        self,
        results: list[dict[str, Any]],
        dsl: dict[str, Any],
        visualization: str,
    ) -> dict[str, Any] | None:
        """Generate a ChartJS configuration from query results.

        Creates a ChartJS-compatible configuration based on the data structure
        and the visualization type returned by the API.

        Args:
            results: Query results as list of dicts
            dsl: The DSL used for the query
            visualization: Visualization type from the API (bar, line, pie, etc.)

        Returns:
            Chart configuration dict or None if not applicable
        """
        if not results or len(results) < 2:
            return None

        # Find label column (first string-typed) and value columns (numeric).
        # Scan up to 5 rows to handle cases where the first row has NULLs.
        columns = list(results[0].keys())
        label_col = None
        value_cols = []

        for col in columns:
            # Find the first non-None value in up to 5 rows
            sample_val = None
            for row in results[:5]:
                v = row.get(col)
                if v is not None:
                    sample_val = v
                    break

            if sample_val is not None and isinstance(sample_val, str) and label_col is None:
                label_col = col
            elif isinstance(sample_val, (int, float, Decimal)):
                value_cols.append(col)

        if not label_col or not value_cols:
            return None

        # Extract data
        labels = [str(row.get(label_col) or "") for row in results]

        colors = [
            "rgba(54, 162, 235, 0.8)",
            "rgba(255, 99, 132, 0.8)",
            "rgba(255, 206, 86, 0.8)",
            "rgba(75, 192, 192, 0.8)",
            "rgba(153, 102, 255, 0.8)",
            "rgba(255, 159, 64, 0.8)",
            "rgba(199, 199, 199, 0.8)",
            "rgba(83, 102, 255, 0.8)",
        ]

        chart_type = visualization if visualization in ("bar", "line", "pie", "doughnut") else "bar"

        datasets = []
        for i, val_col in enumerate(value_cols[:2]):  # Limit to 2 datasets
            values = [float(row.get(val_col) or 0) for row in results]
            dataset: dict[str, Any] = {
                "label": val_col.replace("_", " ").title(),
                "data": values,
                "borderWidth": 1,
            }

            if chart_type in ("pie", "doughnut"):
                # Pie/doughnut: different color per slice
                dataset["backgroundColor"] = colors[: len(labels)]
            elif len(value_cols) == 1:
                # Single dataset: different color per bar
                dataset["backgroundColor"] = colors[: len(labels)]
            else:
                # Multiple datasets: same color per dataset
                dataset["backgroundColor"] = colors[i % len(colors)]
                dataset["borderColor"] = colors[i % len(colors)].replace("0.8", "1")

            datasets.append(dataset)

        title = (
            f"{value_cols[0].replace('_', ' ').title()} by "
            f"{label_col.replace('_', ' ').title()}"
        )

        return {
            "format": "chartjs",
            "title": title,
            "content": {
                "type": chart_type,
                "data": {
                    "labels": labels,
                    "datasets": datasets,
                },
                "options": {
                    "responsive": True,
                    "plugins": {
                        "legend": {"position": "top"},
                        "title": {"display": True, "text": title},
                    },
                },
            },
        }
