"""
Wizard.AI Python Client Library
Official SDK for Wizard.AI API
Created by Arnav Gupta

This client provides a simple interface to interact with Wizard.AI's API,
supporting all personalities (JARVIS, ORACLE, Nerd, Fun, etc.) and both
standard and streaming responses.

Usage:
    from wizard import Wizard
    
    client = Wizard(api_key="your-api-key")
    
    # Simple completion
    response = client.chat.completions.create(
        messages=[
            {"role": "user", "content": "Hello!"}
        ],
        model="wizard-jarvis"
    )
    print(response.choices[0].message.content)
    
    # Streaming
    stream = client.chat.completions.create(
        messages=[{"role": "user", "content": "Tell me a story"}],
        model="wizard-fun",
        stream=True
    )
    for chunk in stream:
        print(chunk.choices[0].delta.content, end="")
"""

import requests
import json
import re
from typing import List, Dict, Optional, Union, Any, Generator
from .errors import AuthenticationError, APIError, RateLimitError, InvalidRequestError

# Version is imported from __init__.py
__all__ = ['Wizard', 'ChatCompletions', 'Completion', 'StreamWrapper']


class Wizard:
    """
    Main Wizard.AI client
    
    This is the main entry point for interacting with the Wizard.AI API.
    Create an instance with your API key and use the `chat` property to
    access chat completion functionality.
    
    Args:
        api_key: Your Wizard.AI API key (starts with 'wiz_')
        base_url: Optional custom base URL (defaults to production API)
        timeout: Request timeout in seconds (default 60)
        validate_key: Whether to validate API key format locally (default True)
    
    Example:
        >>> client = Wizard(api_key="wiz_abc123")
        >>> response = client.chat.completions.create(
        ...     messages=[{"role": "user", "content": "Hello!"}]
        ... )
    """
    
    # API key format validation regex
    API_KEY_PATTERN = re.compile(r'^wiz_[A-Za-z0-9_-]{20,}$')
    
    def __init__(
        self, 
        api_key: str, 
        base_url: str = "https://api.wizardai.dpdns.org/v1",
        timeout: int = 60,
        validate_key: bool = True
    ):
        if not api_key:
            raise ValueError("API key is required")
        
        self.api_key = api_key
        self.base_url = base_url.rstrip('/')
        self.timeout = timeout
        
        # Validate API key format if requested
        if validate_key:
            self._validate_api_key_format()
        
        self._session = requests.Session()
        self._session.headers.update({
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": f"Wizard.AI-Python-SDK/{self._get_version()}"
        })
        
        # Initialize endpoints
        self.chat = ChatCompletions(self)
    
    def _validate_api_key_format(self):
        """Validate API key format locally before sending to server"""
        if not self.API_KEY_PATTERN.match(self.api_key):
            raise AuthenticationError(
                "Invalid API key format. "
                "API keys should start with 'wiz_' and be at least 20 characters long. "
                "Get your API key from: https://wizardai.dpdns.org/api-keys"
            )
    
    def _get_version(self) -> str:
        """Get package version"""
        try:
            from . import __version__
            return __version__
        except ImportError:
            return "unknown"
    
    def _request(
        self,
        method: str,
        endpoint: str,
        json_data: Optional[Dict] = None,
        stream: bool = False
    ) -> Union[Dict, Generator]:
        """
        Make HTTP request to Wizard.AI API
        
        Args:
            method: HTTP method (GET, POST, etc.)
            endpoint: API endpoint (e.g., "chat/completions")
            json_data: Request body as dictionary
            stream: Whether to stream the response
            
        Returns:
            Parsed JSON response or generator for streaming
        """
        url = f"{self.base_url}/{endpoint.lstrip('/')}"
        
        if stream:
            return self._stream_request(method, url, json_data)
        
        try:
            response = self._session.request(
                method=method,
                url=url,
                json=json_data,
                timeout=self.timeout
            )
        except requests.exceptions.ConnectionError:
            raise APIError("Failed to connect to Wizard.AI API. Check your internet connection.")
        except requests.exceptions.Timeout:
            raise APIError(f"Request timed out after {self.timeout} seconds")
        except requests.exceptions.RequestException as e:
            raise APIError(f"Request failed: {str(e)}")
        
        if response.status_code == 200:
            return response.json()
        elif response.status_code == 401:
            # Server-side key validation failed
            raise AuthenticationError(
                "Invalid API key. Please check your credentials. "
                "Get your API key from: https://wizardai.dpdns.org/api-keys"
            )
        elif response.status_code == 403:
            raise AuthenticationError("API key does not have permission for this action.")
        elif response.status_code == 404:
            raise APIError(f"Endpoint not found: {endpoint}")
        elif response.status_code == 429:
            raise RateLimitError("Rate limit exceeded. Please try again later.")
        elif response.status_code >= 500:
            raise APIError(f"Wizard.AI server error (HTTP {response.status_code}). Please try again later.")
        else:
            self._handle_error_response(response)
    
    def _stream_request(self, method: str, url: str, json_data: Optional[Dict] = None) -> Generator:
        """
        Handle streaming requests with Server-Sent Events
        
        Args:
            method: HTTP method
            url: Full URL
            json_data: Request body
            
        Yields:
            Parsed JSON chunks from the stream
        """
        try:
            response = self._session.request(
                method=method,
                url=url,
                json=json_data,
                stream=True,
                timeout=self.timeout
            )
        except requests.exceptions.RequestException as e:
            raise APIError(f"Streaming request failed: {str(e)}")
        
        if response.status_code != 200:
            if response.status_code == 401:
                raise AuthenticationError("Invalid API key")
            elif response.status_code == 429:
                raise RateLimitError("Rate limit exceeded")
            elif response.status_code >= 500:
                raise APIError(f"Server error (HTTP {response.status_code})")
            else:
                self._handle_error_response(response)
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    data = line[6:]  # Remove 'data: ' prefix
                    if data == '[DONE]':
                        break
                    try:
                        yield json.loads(data)
                    except json.JSONDecodeError:
                        # Ignore malformed JSON
                        continue
    
    def _handle_error_response(self, response):
        """Parse and raise appropriate error from response"""
        try:
            error_data = response.json()
            error_msg = error_data.get('error', {}).get('message', 'Unknown error')
            if 'type' in error_data.get('error', {}):
                error_type = error_data['error']['type']
                if 'rate_limit' in error_type.lower():
                    raise RateLimitError(error_msg)
                elif 'auth' in error_type.lower():
                    raise AuthenticationError(error_msg)
        except (json.JSONDecodeError, AttributeError):
            error_msg = f"HTTP {response.status_code}: {response.reason}"
        
        raise APIError(f"Wizard.AI API Error: {error_msg}")
    
    def list_models(self) -> List[Dict]:
        """
        List all available models/personalities
        
        Returns:
            List of model objects with id, name, and capabilities
        
        Example:
            >>> client.list_models()
            [
                {'id': 'wizard-jarvis', 'object': 'model', ...},
                {'id': 'wizard-oracle', ...}
            ]
        """
        return self._request("GET", "models").get('data', [])
    
    def validate_key(self) -> bool:
        """
        Validate the API key by making a test request
        
        Returns:
            True if key is valid, raises exception otherwise
        
        Example:
            >>> if client.validate_key():
            ...     print("✅ API key is valid!")
        """
        try:
            self.list_models()
            return True
        except AuthenticationError:
            return False
    
    def close(self):
        """Close the underlying session"""
        self._session.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class ChatCompletions:
    """
    Chat completions endpoint handler
    
    This class provides methods for creating chat completions with
    Wizard.AI's various personalities.
    """
    
    def __init__(self, client: Wizard):
        self.client = client
    
    def create(
        self,
        messages: List[Dict[str, str]],
        model: str = "wizard-ai-pro",
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        top_p: float = 1.0,
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        stop: Optional[Union[str, List[str]]] = None,
        stream: bool = False,
        **kwargs
    ) -> Union['Completion', 'StreamWrapper']:
        """
        Create a chat completion
        
        This is the main method for interacting with Wizard.AI's chat models.
        It supports both standard and streaming responses.
        
        Args:
            messages: List of message objects with 'role' and 'content'
                Each message can have role: 'system', 'user', or 'assistant'
            model: Model/personality to use. Options:
                - wizard-jarvis: Sophisticated AI assistant
                - wizard-oracle: Mystical, all-knowing
                - wizard-nerd: Detailed, academic
                - wizard-fun: Playful, energetic
                - wizard-sarcastic: Witty, sarcastic
                - wizard-fast: Quick responses
                - wizard-normal: Balanced conversation
                - wizard-ai-pro: Default (JARVIS)
            temperature: Sampling temperature (0-2). Higher values = more random
            max_tokens: Maximum number of tokens to generate
            top_p: Nucleus sampling parameter (0-1)
            frequency_penalty: Penalize frequent tokens (-2 to 2)
            presence_penalty: Penalize tokens that have appeared (-2 to 2)
            stop: Stop sequence(s) that end generation
            stream: Whether to stream the response token by token
            **kwargs: Additional parameters to pass to the API
        
        Returns:
            If stream=False: A Completion object with the full response
            If stream=True: A StreamWrapper that yields StreamChunk objects
        
        Raises:
            AuthenticationError: Invalid API key
            RateLimitError: Rate limit exceeded
            APIError: Other API errors
        
        Examples:
            >>> # Standard completion
            >>> response = client.chat.completions.create(
            ...     messages=[
            ...         {"role": "system", "content": "You are JARVIS"},
            ...         {"role": "user", "content": "What's the status?"}
            ...     ],
            ...     model="wizard-jarvis"
            ... )
            >>> print(response.choices[0].message.content)
            
            >>> # Streaming
            >>> stream = client.chat.completions.create(
            ...     messages=[{"role": "user", "content": "Tell a story"}],
            ...     model="wizard-fun",
            ...     stream=True
            ... )
            >>> for chunk in stream:
            ...     print(chunk.choices[0].delta.content, end="")
        """
        # Validate messages
        if not messages or not isinstance(messages, list):
            raise InvalidRequestError("Messages must be a non-empty list")
        
        for msg in messages:
            if not isinstance(msg, dict):
                raise InvalidRequestError("Each message must be a dictionary")
            if 'role' not in msg or 'content' not in msg:
                raise InvalidRequestError("Each message must have 'role' and 'content' keys")
            if msg['role'] not in ['system', 'user', 'assistant']:
                raise InvalidRequestError(f"Invalid role: {msg['role']}. Must be 'system', 'user', or 'assistant'")
        
        # Validate temperature
        if not 0 <= temperature <= 2:
            raise InvalidRequestError("Temperature must be between 0 and 2")
        
        # Validate penalties
        if not -2 <= frequency_penalty <= 2:
            raise InvalidRequestError("Frequency penalty must be between -2 and 2")
        if not -2 <= presence_penalty <= 2:
            raise InvalidRequestError("Presence penalty must be between -2 and 2")
        
        # Prepare request data
        data = {
            "messages": messages,
            "model": model,
            "temperature": temperature,
            "top_p": top_p,
            "frequency_penalty": frequency_penalty,
            "presence_penalty": presence_penalty,
            "stream": stream,
            **kwargs
        }
        
        if max_tokens is not None:
            if max_tokens <= 0:
                raise InvalidRequestError("max_tokens must be positive")
            data["max_tokens"] = max_tokens
        
        if stop is not None:
            data["stop"] = stop
        
        response = self.client._request(
            "POST", 
            "chat/completions", 
            json_data=data,
            stream=stream
        )
        
        if stream:
            return StreamWrapper(response)
        
        return Completion(response)
    
    # Alias for backward compatibility
    completions = create


class Completion:
    """
    Wrapper for non-streaming completion response
    
    Provides easy access to the completion results and metadata.
    """
    
    def __init__(self, data: Dict):
        self._data = data
        self.id = data.get('id')
        self.object = data.get('object')
        self.created = data.get('created')
        self.model = data.get('model')
        self.choices = [Choice(c) for c in data.get('choices', [])]
        self.usage = Usage(data.get('usage', {}))
    
    def __repr__(self):
        return f"<Completion id={self.id} model={self.model} tokens={self.usage.total_tokens}>"
    
    def __str__(self):
        if self.choices:
            return self.choices[0].message.content
        return ""


class Choice:
    """Wrapper for a single completion choice"""
    
    def __init__(self, data: Dict):
        self.index = data.get('index')
        self.message = Message(data.get('message', {}))
        self.finish_reason = data.get('finish_reason')
    
    def __repr__(self):
        return f"<Choice index={self.index} finish_reason={self.finish_reason}>"


class Message:
    """Wrapper for a message object in non-streaming responses"""
    
    def __init__(self, data: Dict):
        self.role = data.get('role')
        self.content = data.get('content')
    
    def __repr__(self):
        return f"<Message role={self.role}>"
    
    def __str__(self):
        return self.content or ""


class Delta:
    """Wrapper for streaming delta content"""
    
    def __init__(self, data: Dict):
        self.role = data.get('role')
        self.content = data.get('content')
    
    def __repr__(self):
        return f"<Delta content={self.content}>"


class StreamChoice:
    """Wrapper for a streaming choice containing delta updates"""
    
    def __init__(self, data: Dict):
        self.index = data.get('index')
        self.delta = Delta(data.get('delta', {}))
        self.finish_reason = data.get('finish_reason')
    
    def __repr__(self):
        return f"<StreamChoice index={self.index}>"


class StreamWrapper:
    """
    Wrapper for streaming responses
    
    This class allows iteration over streaming chunks from the API.
    Each iteration yields a StreamChunk object.
    """
    
    def __init__(self, generator):
        self.generator = generator
    
    def __iter__(self):
        return self
    
    def __next__(self):
        try:
            chunk = next(self.generator)
            return StreamChunk(chunk)
        except StopIteration:
            raise StopIteration
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


class StreamChunk:
    """Wrapper for a single streaming chunk"""
    
    def __init__(self, data: Dict):
        self._data = data
        self.id = data.get('id')
        self.object = data.get('object')
        self.created = data.get('created')
        self.model = data.get('model')
        self.choices = [StreamChoice(c) for c in data.get('choices', [])]
    
    def __repr__(self):
        return f"<StreamChunk id={self.id}>"


class Usage:
    """Wrapper for token usage information"""
    
    def __init__(self, data: Dict):
        self.prompt_tokens = data.get('prompt_tokens', 0)
        self.completion_tokens = data.get('completion_tokens', 0)
        self.total_tokens = data.get('total_tokens', 0)
    
    def __repr__(self):
        return f"<Usage total_tokens={self.total_tokens}>"