"""
API client for PCompressLR cloud service.

SDK v0.2.0 notes:
- `complete()` is messages-first (defaults to async job endpoints; `mode="sync"` uses POST /api/v2/complete)
- Provider API keys are intentionally not accepted by the SDK (BYOK via dashboard)
"""

import os
import requests
import time
from typing import Dict, Optional, Literal, List, Any
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from uuid import uuid4

# Try to load .env file if python-dotenv is available
try:
    from dotenv import load_dotenv
    from pathlib import Path
    
    # Try to find .env file in multiple locations:
    # 1. Backend directory (where this package is located)
    # 2. Current working directory
    # 3. backend/.env from project root
    
    # Get the directory where this file is located
    # Path(__file__) = backend/pcompresslr/api_client.py
    # .parent = backend/pcompresslr/
    # .parent = backend/
    current_file_dir = Path(__file__).parent.parent
    
    # Try loading from backend directory first
    env_paths = [
        current_file_dir / ".env",  # backend/.env
        Path.cwd() / ".env",  # Current working directory
        Path.cwd() / "backend" / ".env",  # backend/.env from project root
    ]
    
    # Try each path
    loaded = False
    for env_path in env_paths:
        if env_path.exists():
            load_dotenv(env_path, override=True)
            loaded = True
            break
    
    # If no .env file found in specific locations, use default behavior (searches upward)
    if not loaded:
        load_dotenv()
except ImportError:
    # python-dotenv not installed, skip .env loading
    pass
except Exception:
    # If anything goes wrong, just continue without .env loading
    pass


class PcompresslrAPIError(Exception):
    """Base exception for API errors."""
    pass


class APIKeyError(PcompresslrAPIError):
    """Exception raised when API key is invalid or missing."""
    pass


class RateLimitError(PcompresslrAPIError):
    """Exception raised when rate limit is exceeded."""
    pass


class APIRequestError(PcompresslrAPIError):
    """Exception raised for general API request errors."""
    pass


class PcompresslrAPIClient:
    """
    Client for interacting with PCompressLR cloud API.
    
    Usage:
        client = PcompresslrAPIClient(api_key="your-key")
        result = client.compress("your prompt", model="gpt-4")
    """
    
    DEFAULT_API_URL = "https://api.compress.lightreach.io"
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
        timeout: int = 900  # 15 minutes - complete() can include long upstream LLM calls
    ):
        """
        Initialize the API client.
        
        Args:
            api_key: API key for authentication. If not provided, will check
                    PCOMPRESLR_API_KEY environment variable.
            api_url: Base URL for the API. If not provided, uses default production URL.
                    Can also be set via PCOMPRESLR_API_URL environment variable.
                    This is primarily for internal/testing use - users should use the default.
            timeout: Request timeout in seconds
        
        Raises:
            APIKeyError: If API key is not provided and not found in environment
        """
        # Get API key from parameter or environment
        # Support both legacy env var and the more semantic name.
        self.api_key = api_key or os.getenv("LIGHTREACH_API_KEY") or os.getenv("PCOMPRESLR_API_KEY")
        if not self.api_key:
            raise APIKeyError(
                "API key is required. Provide it as a parameter or set "
                "PCOMPRESLR_API_KEY environment variable."
            )
        
        # Basic validation - API keys should not be empty strings
        if not self.api_key.strip():
            raise APIKeyError(
                "API key cannot be empty. Please provide a valid API key."
            )
        
        # Get API URL from parameter or environment or use default
        self.api_url = (api_url or 
                       os.getenv("PCOMPRESLR_API_URL") or 
                       self.DEFAULT_API_URL).rstrip('/')
        
        self.timeout = timeout
        
        # Create session with retry strategy
        # Note: We don't retry on 401 (invalid API key) or 4xx errors - fail fast
        self.session = requests.Session()
        retry_strategy = Retry(
            total=2,  # Reduced from 3 for faster failure
            backoff_factor=0.3,  # Reduced backoff for faster retries
            status_forcelist=[429, 500, 502, 503, 504],  # Only retry on these status codes
            allowed_methods=["POST"],
            respect_retry_after_header=True
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)
        
        # Set default headers
        try:
            from ._version import __version__ as _sdk_version
        except Exception:
            _sdk_version = "0.2.0"
        self.session.headers.update(
            {
                # Prefer standard Bearer auth, but keep X-API-Key for backward compatibility.
                "Authorization": f"Bearer {self.api_key}",
                "X-API-Key": self.api_key,
                "Content-Type": "application/json",
                "User-Agent": f"compress-lightreach-python/{_sdk_version}",
            }
        )
    
    def _make_request(
        self,
        endpoint: str,
        data: Dict,
        method: str = "POST",
        *,
        timeout: Optional[float] = None,
        headers: Optional[Dict[str, str]] = None,
    ) -> Dict:
        """
        Make an HTTP request to the API.
        
        Args:
            endpoint: API endpoint path (e.g., "/api/v1/compress")
            data: Request payload
            method: HTTP method (default: "POST")
        
        Returns:
            JSON response as dictionary
        
        Raises:
            APIKeyError: If API key is invalid
            RateLimitError: If rate limit is exceeded
            APIRequestError: For other API errors
        """
        url = f"{self.api_url}{endpoint}"
        
        try:
            req_timeout = self.timeout if timeout is None else timeout
            req_headers = dict(self.session.headers)
            if headers:
                req_headers.update(headers)

            kwargs: Dict[str, Any] = {"timeout": req_timeout, "headers": req_headers}
            if method.upper() == "GET":
                kwargs["params"] = data
            else:
                kwargs["json"] = data

            response = self.session.request(method=method, url=url, **kwargs)

            def _safe_body_snippet(max_len: int = 800) -> str:
                try:
                    t = response.text or ""
                except Exception:
                    t = ""
                t = t.strip()
                if not t:
                    return ""
                if len(t) <= max_len:
                    return t
                return t[:max_len] + " …(truncated)…"
            
            # Handle different status codes
            if response.status_code == 401:
                try:
                    error_detail = response.json().get("detail", "Invalid API key")
                except:
                    error_detail = "Invalid API key"
                raise APIKeyError(
                    f"Authentication failed: {error_detail}. "
                    "Please check your API key at https://compress.lightreach.io"
                )
            
            elif response.status_code == 429:
                try:
                    error_detail = response.json().get("detail", "Rate limit exceeded")
                except Exception:
                    error_detail = _safe_body_snippet() or "Rate limit exceeded"
                raise RateLimitError(f"Rate limit exceeded: {error_detail}")
            
            elif response.status_code == 400:
                try:
                    error_detail = response.json().get("detail", "Bad request")
                except Exception:
                    error_detail = _safe_body_snippet() or "Bad request"
                raise APIRequestError(f"Bad request: {error_detail}")
            
            elif response.status_code == 403:
                try:
                    error_detail = response.json().get("detail", "Forbidden")
                except Exception:
                    error_detail = _safe_body_snippet() or "Forbidden"
                raise APIRequestError(f"Forbidden: {error_detail}")
            
            elif not response.ok:
                try:
                    error_detail = response.json().get("detail", response.text)
                except:
                    error_detail = _safe_body_snippet() or f"HTTP {response.status_code}"
                raise APIRequestError(
                    f"API request failed (HTTP {response.status_code}): {error_detail}"
                )

            # Success path: some proxies/edge cases return empty or non-JSON bodies.
            try:
                return response.json()
            except Exception:
                snippet = _safe_body_snippet()
                raise APIRequestError(
                    f"API returned a non-JSON response (HTTP {response.status_code}). "
                    f"URL={url}. Body={snippet or '(empty)'}"
                )
        
        except requests.exceptions.Timeout:
            raise APIRequestError(
                f"Request to {url} timed out after {self.timeout} seconds. "
                "The API may be experiencing high load or may be unreachable. "
                "Please check your internet connection and try again."
            )
        
        except requests.exceptions.ConnectionError as e:
            error_msg = str(e)
            # Check if this is a timeout-related connection error
            if "timed out" in error_msg.lower() or "timeout" in error_msg.lower():
                raise APIRequestError(
                    f"Failed to connect to API at {self.api_url}. "
                    f"Error: {error_msg}. "
                    "Please check your internet connection and verify the API endpoint is accessible.\n\n"
                    f"If this problem persists, please check https://compress.lightreach.io/status or contact support."
                )
            raise APIRequestError(
                f"Failed to connect to API at {self.api_url}. "
                f"Error: {error_msg}. "
                "Please check your internet connection and verify the API endpoint is accessible.\n\n"
                f"If this problem persists, please check https://compress.lightreach.io/status or contact support."
            )
        
        except (APIKeyError, RateLimitError, APIRequestError):
            # Re-raise our custom exceptions
            raise
        
        except requests.exceptions.RequestException as e:
            raise APIRequestError(f"Request failed: {str(e)}")

    @staticmethod
    def _to_openai_superset_response(raw: Dict[str, Any]) -> Dict[str, Any]:
        """
        Normalize /api/v2/complete-style payload into an OpenAI-style superset.

        Returns a dict that always contains standard chat completion fields:
        id, object, created, model, choices, usage
        while preserving all LightReach-native fields plus a `lightreach` extension block.
        """
        if not isinstance(raw, dict):
            return raw

        llm_stats = raw.get("llm_stats") if isinstance(raw.get("llm_stats"), dict) else {}
        routing_info = raw.get("routing_info") if isinstance(raw.get("routing_info"), dict) else {}

        content = (
            raw.get("content")
            or raw.get("decompressed_response")
            or raw.get("text")
            or ""
        )
        if not isinstance(content, str):
            content = str(content)

        model_name = (
            raw.get("model")
            or raw.get("model_used")
            or routing_info.get("selected_model")
            or llm_stats.get("model")
            or "lightreach"
        )
        finish_reason = llm_stats.get("finish_reason") or "stop"
        prompt_tokens = int(llm_stats.get("input_tokens") or 0)
        completion_tokens = int(llm_stats.get("output_tokens") or 0)
        total_tokens = int(llm_stats.get("total_tokens") or (prompt_tokens + completion_tokens))

        message: Dict[str, Any] = {"role": "assistant", "content": content}
        tool_calls = raw.get("tool_calls")
        if isinstance(tool_calls, list) and tool_calls:
            message["tool_calls"] = tool_calls
            if not content:
                message["content"] = None

        openai_fields: Dict[str, Any] = {
            "id": str(raw.get("id") or f"chatcmpl-{uuid4().hex}"),
            "object": str(raw.get("object") or "chat.completion"),
            "created": int(raw.get("created") or int(time.time())),
            "model": str(model_name),
            "choices": raw.get("choices")
            if isinstance(raw.get("choices"), list)
            else [
                {
                    "index": 0,
                    "message": message,
                    "finish_reason": finish_reason,
                }
            ],
            "usage": raw.get("usage")
            if isinstance(raw.get("usage"), dict)
            else {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": total_tokens,
            },
        }

        out: Dict[str, Any] = dict(raw)
        for k, v in openai_fields.items():
            out[k] = v

        out.setdefault("content", content)
        out.setdefault("decompressed_response", content)
        out.setdefault("text", content)

        out["lightreach"] = {
            "content": content,
            "compression_stats": out.get("compression_stats"),
            "llm_stats": out.get("llm_stats"),
            "routing_info": out.get("routing_info"),
            "warnings": out.get("warnings"),
            "tokens_saved": out.get("tokens_saved"),
            "tokens_used": out.get("tokens_used"),
            "compression_ratio": out.get("compression_ratio"),
            "cost_estimate": out.get("cost_estimate"),
            "savings_estimate": out.get("savings_estimate"),
            "provider_used": out.get("provider_used"),
            "model_used": out.get("model_used"),
            "latency_ms": (
                out.get("llm_stats", {}).get("latency_ms")
                if isinstance(out.get("llm_stats"), dict)
                else None
            ),
        }
        return out

    # ==================== Async /complete jobs (Option C) ====================

    def create_complete_job(
        self,
        data: Dict[str, Any],
        *,
        idempotency_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Create an async /complete job (POST /api/v1/complete/jobs).
        Returns: {job_id, status, status_url}
        """
        extra_headers: Dict[str, str] = {}
        if idempotency_key:
            extra_headers["Idempotency-Key"] = idempotency_key
        return self._make_request("/api/v1/complete/jobs", data, headers=extra_headers)

    def get_complete_job(self, job_id: str) -> Dict[str, Any]:
        """
        Poll async job status (GET /api/v1/complete/jobs/{job_id}).
        Returns: {job_id, status, phase, progress, result?, error?}
        """
        if not job_id:
            raise ValueError("job_id is required")
        # Keep polling calls short; this avoids client-side long hangs.
        return self._make_request(f"/api/v1/complete/jobs/{job_id}", data={}, method="GET", timeout=min(30, float(self.timeout)))

    def wait_for_complete_job(
        self,
        job_id: str,
        *,
        poll_interval_s: float = 1.0,
        max_wait_s: Optional[float] = None,
    ) -> Dict[str, Any]:
        """
        Wait until a job completes and return the final result dict.
        """
        if max_wait_s is None:
            max_wait_s = float(self.timeout)
        deadline = time.time() + float(max_wait_s)

        interval = max(0.2, float(poll_interval_s))
        while True:
            status = self.get_complete_job(job_id)
            s = (status.get("status") or "").strip().lower()
            if s == "succeeded":
                result = status.get("result")
                if isinstance(result, dict):
                    return result
                raise APIRequestError("Async job succeeded but result payload was missing.")
            if s in ("failed", "canceled"):
                raise APIRequestError(f"Async complete job {s}: {status.get('error')}")

            if time.time() > deadline:
                raise APIRequestError(
                    f"Async complete job timed out after {max_wait_s}s (status={status.get('status')})."
                )

            time.sleep(interval)
            interval = min(interval * 1.2, 2.0)
    
    def compress(
        self,
        prompt: str,
        model: str = "gpt-4",
        tags: Optional[Dict[str, str]] = None
    ) -> Dict:
        """
        Compress a prompt using the cloud API.
        
        Args:
            prompt: The prompt text to compress
            model: LLM model name for tokenization (e.g., 'gpt-4', 'gpt-3.5-turbo')
            tags: Optional tags for cost attribution. Supported keys:
                  'team', 'environment', 'feature' (e.g., {'team': 'backend', 'environment': 'production'}).
                  Values are validated server-side against your workspace's allowed list.
        
        Returns:
            Dictionary containing:
                - compressed: Compressed prompt text
                - dictionary: Decompression dictionary
                - llm_format: LLM-ready formatted string
                - compression_ratio: Compression ratio
                - original_size: Original size in characters
                - compressed_size: Compressed size in characters
                - processing_time_ms: Processing time in milliseconds
                - algorithm: Algorithm used
        
        Raises:
            APIKeyError: If API key is invalid
            RateLimitError: If rate limit is exceeded
            APIRequestError: For other API errors
        """
        data = {
            "prompt": prompt,
            "model": model,
            "algorithm": "greedy"
        }
        
        if tags is not None:
            data["tags"] = tags
        
        return self._make_request("/api/v1/compress", data)
    
    def decompress(self, llm_format: str) -> Dict:
        """
        Decompress an LLM-formatted compressed prompt.
        
        Args:
            llm_format: LLM-formatted compressed prompt string
        
        Returns:
            Dictionary containing:
                - decompressed: Decompressed original prompt
                - processing_time_ms: Processing time in milliseconds
        
        Raises:
            APIKeyError: If API key is invalid
            RateLimitError: If rate limit is exceeded
            APIRequestError: For other API errors
        """
        data = {
            "llm_format": llm_format
        }
        
        return self._make_request("/api/v1/decompress", data)
    
    def health_check(self) -> Dict:
        """
        Check API health status.
        
        Returns:
            Dictionary with health status information
        
        Raises:
            APIRequestError: If health check fails
        """
        try:
            response = self.session.get(
                f"{self.api_url}/health",
                timeout=5
            )
            if response.ok:
                return response.json()
            else:
                raise APIRequestError(f"Health check failed: HTTP {response.status_code}")
        except requests.exceptions.RequestException as e:
            raise APIRequestError(f"Health check failed: {str(e)}")
    
    def complete(
        self,
        messages: List[Dict[str, Any]],
        llm_provider: Optional[Literal["openai", "anthropic", "google", "deepseek", "moonshot"]] = None,
        desired_hle: Optional[float] = None,
        compress: bool = True,
        compression_config: Optional[Dict[str, Any]] = None,
        compress_output: bool = False,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        tags: Optional[Dict[str, str]] = None,
        max_history_messages: Optional[int] = None,
        *,
        mode: Literal["async", "sync"] = "async",
        wait: bool = True,
        poll_interval_s: float = 1.0,
        max_wait_s: Optional[float] = None,
        idempotency_key: Optional[str] = None,
        # Deprecated parameters (v0.2.0, ignored in v1.0.0)
        model: Optional[str] = None,
        hle_target_percent: Optional[float] = None,
        min_hle_score: Optional[float] = None,
        auto_select_by_hle: Optional[bool] = None,
        same_provider_only: Optional[bool] = None,
    ) -> Dict:
        """
        Complete a conversation with intelligent model selection, compression, LLM call, and decompression.

        Default (mode="async"):
        - Enqueues a job via POST /api/v1/complete/jobs
        - Polls GET /api/v1/complete/jobs/{job_id} until completion (unless wait=False)

        Legacy (mode="sync"):
        - Calls POST /api/v2/complete and waits for the full response (can time out client-side)

        Model routing is managed by your workspace admin via the dashboard. The system
        selects the cheapest model that meets the admin-configured quality ceiling
        (global, tag-level, or integration-level). HLE scores are maintained server-side.

        Provider API keys must be stored in your account (BYOK via dashboard).
        
        Args:
            messages: Full conversation history as list of {"role": "...", "content": "..."} objects (required)
            llm_provider: Optional provider constraint (e.g., 'openai', 'anthropic', 'google').
                         Omit to allow cross-provider cost optimization. (default: None)
            compress: Whether to compress eligible messages (default: True)
            compression_config: Optional per-role compression configuration
            compress_output: Whether to request compressed output from LLM (default: False)
            temperature: LLM temperature setting (optional)
            max_tokens: Maximum tokens to generate (optional)
            tags: Tags for cost attribution and tag-level quality ceilings. Supported keys:
                  'team', 'environment', 'feature' (e.g., {'team': 'backend', 'environment': 'production'}).
                  Values are validated server-side against your workspace's allowed list.
                  The 'integration' tag is reserved for system use and should not be set manually.
            max_history_messages: Optional limit on conversation history length
        
        Returns:
            Dictionary containing:
                - id/object/created/model/choices/usage: OpenAI-compatible chat completion fields
                - content: Final response text (alias of choices[0].message.content)
                - compression_stats: Input compression statistics
                - llm_stats: LLM usage statistics
                - routing_info: Detailed routing decision:
                    - selected_model: Model chosen by system
                    - selected_provider: Provider chosen by system
                    - model_hle: HLE score of selected model (server-computed)
                    - model_price_per_million: Blended price per million tokens
                    - effective_hle: The quality ceiling that was applied
                    - hle_source: "tag", "global", or "none"
                - warnings: Any warnings (tag normalization, etc.)
                - cost_estimate: Estimated USD cost
                - savings_estimate: Estimated USD savings
        
        Raises:
            APIKeyError: If API key is invalid
            RateLimitError: If rate limit is exceeded
            APIRequestError: For configuration errors (e.g., no provider keys, invalid tags)
        
        Example:
            response = client.complete(
                messages=[{"role": "user", "content": "Hello"}],
                tags={"team": "backend", "environment": "production"},
            )
            
            print(f"Selected: {response['routing_info']['selected_model']}")
            print(f"HLE ceiling: {response['routing_info']['effective_hle']}")
        """
        if not messages or not isinstance(messages, list):
            raise ValueError("messages is required and must be a non-empty list")

        # Warn about deprecated parameters
        if model is not None:
            import warnings
            warnings.warn(
                "Parameter 'model' is deprecated in v1.0.0 and will be ignored. "
                "The system now selects models automatically based on your provider keys and HLE requirements.",
                DeprecationWarning,
                stacklevel=2
            )
        
        if any(p is not None for p in [hle_target_percent, min_hle_score, auto_select_by_hle, same_provider_only]):
            import warnings
            warnings.warn(
                "Parameters 'hle_target_percent', 'min_hle_score', 'auto_select_by_hle', and 'same_provider_only' "
                "are deprecated in v1.0.0. Use 'desired_hle' and optional 'llm_provider' instead.",
                DeprecationWarning,
                stacklevel=2
            )

        data: Dict[str, Any] = {
            "messages": messages,
            "compress": compress,
            "compress_output": compress_output,
            "algorithm": "greedy",
        }

        # v1.0.0 parameters
        if llm_provider is not None:
            data["llm_provider"] = llm_provider
        
        if desired_hle is not None:
            data["desired_hle"] = desired_hle

        if compression_config is not None:
            data["compression_config"] = compression_config

        if max_history_messages is not None:
            data["max_history_messages"] = max_history_messages
        
        if temperature is not None:
            data["temperature"] = temperature
        if max_tokens is not None:
            data["max_tokens"] = max_tokens
        if tags is not None:
            data["tags"] = tags

        if mode == "sync":
            return self._to_openai_superset_response(self._make_request("/api/v2/complete", data))

        idem = (idempotency_key or "").strip() or f"sdk-{uuid4()}"
        job = self.create_complete_job(data, idempotency_key=idem)
        if not wait:
            return job
        job_id = str(job.get("job_id") or "")
        if not job_id:
            raise APIRequestError(f"Async job creation failed: {job}")
        result = self.wait_for_complete_job(job_id, poll_interval_s=poll_interval_s, max_wait_s=max_wait_s)
        return self._to_openai_superset_response(result)
