"""Query runner for executing BTQL queries against Braintrust API."""

import asyncio
import gzip
import json
import re
import httpx

from .models import QueryResult, QueryStatus


# Braintrust org endpoint mappings
BRAINTRUST_ORG_BTQL_ENDPOINTS = {
    "bt-azure": "https://prod-msft-api-endpoint-eackahcmezgzaybz.a03.azurefd.net/btql",
    "non-pme": "https://api.braintrust.dev/btql",
}


class QueryRunner:
    """Executes BTQL queries against the Braintrust API."""

    MAX_RETRIES = 3

    def __init__(self, api_key: str, timeout: int = 60, braintrust_org: str = "non-pme") -> None:
        self.api_key = api_key
        self.timeout = timeout
        self.braintrust_org = braintrust_org
        self.BTQL_ENDPOINT = BRAINTRUST_ORG_BTQL_ENDPOINTS.get(braintrust_org, BRAINTRUST_ORG_BTQL_ENDPOINTS["non-pme"])

    async def execute(self, query: str) -> QueryResult:
        """Execute a BTQL query and return the result with retry for rate limits."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {"query": query}

        for retry in range(self.MAX_RETRIES):
            try:
                async with httpx.AsyncClient(timeout=self.timeout) as client:
                    response = await client.post(
                        self.BTQL_ENDPOINT,
                        json=payload,
                        headers=headers,
                    )

                    if response.status_code == 200:
                        response_json = response.json()
                        data = response_json.get("data", [])
                        return QueryResult(
                            status=QueryStatus.SUCCESS,
                            data=data,
                            raw_response=response_json,
                        )
                    elif response.status_code == 303:
                        # Handle redirect to S3 for large responses
                        redirect_url = response.headers.get("Location")
                        if not redirect_url:
                            # Try to get from response body
                            try:
                                body = response.json()
                                redirect_url = body.get("url") or body.get("location")
                            except Exception:
                                pass

                        if redirect_url:
                            return await self._fetch_redirect(client, redirect_url)
                        else:
                            return QueryResult(
                                status=QueryStatus.ERROR,
                                error="HTTP 303: Redirect without Location header",
                            )
                    elif response.status_code == 429:
                        # Rate limited - extract retry time and wait
                        retry_after = self._extract_retry_after(response.text)
                        if retry < self.MAX_RETRIES - 1:
                            await asyncio.sleep(retry_after)
                            continue
                        else:
                            return QueryResult(
                                status=QueryStatus.ERROR,
                                error=f"HTTP 429: Rate limited. Retry after {retry_after:.1f}s",
                            )
                    else:
                        error_text = response.text
                        try:
                            error_json = response.json()
                            if "error" in error_json:
                                error_text = error_json["error"]
                            elif "message" in error_json:
                                error_text = error_json["message"]
                        except Exception:
                            pass

                        return QueryResult(
                            status=QueryStatus.ERROR,
                            error=f"HTTP {response.status_code}: {error_text}",
                            raw_response={"status_code": response.status_code, "body": response.text},
                        )

            except httpx.TimeoutException:
                return QueryResult(
                    status=QueryStatus.ERROR,
                    error=f"Request timed out after {self.timeout} seconds",
                )
            except httpx.RequestError as e:
                return QueryResult(
                    status=QueryStatus.ERROR,
                    error=f"Request failed: {str(e)}",
                )

        return QueryResult(
            status=QueryStatus.ERROR,
            error="Max retries exceeded",
        )

    def _extract_retry_after(self, response_text: str) -> float:
        """Extract retry-after time from error response."""
        try:
            # Try to parse JSON and get retry time
            match = re.search(r'Retry after ([\d.]+) seconds', response_text)
            if match:
                return float(match.group(1)) + 1  # Add 1 second buffer
        except Exception:
            pass
        return 10.0  # Default to 10 seconds

    async def _fetch_redirect(self, client: httpx.AsyncClient, url: str) -> QueryResult:
        """Fetch data from a redirect URL (typically S3)."""
        try:
            response = await client.get(url)
            if response.status_code == 200:
                # Check if gzipped
                content = response.content
                if url.endswith('.gz') or response.headers.get('Content-Encoding') == 'gzip':
                    try:
                        content = gzip.decompress(content)
                    except Exception:
                        pass  # Not gzipped or decompression failed

                response_json = json.loads(content)
                data = response_json.get("data", [])
                return QueryResult(
                    status=QueryStatus.SUCCESS,
                    data=data,
                    raw_response=response_json,
                )
            else:
                return QueryResult(
                    status=QueryStatus.ERROR,
                    error=f"Redirect fetch failed: HTTP {response.status_code}",
                )
        except Exception as e:
            return QueryResult(
                status=QueryStatus.ERROR,
                error=f"Redirect fetch failed: {str(e)}",
            )
