"""Main Mia21 client implementation."""

import requests
import json
import uuid
from typing import Optional, List, Generator, Dict, Any
from .models import ChatMessage, Space, InitializeResponse, ChatResponse
from .exceptions import Mia21Error, ChatNotInitializedError, APIError


class Mia21Client:
    """
    Mia21 Chat API Client
    
    Example:
        >>> from mia21 import Mia21Client
        >>> client = Mia21Client(api_key="your-api-key")
        >>> client.initialize()
        >>> response = client.chat("Hello!")
        >>> print(response.message)
    """
    
    def __init__(
        self,
        api_key: str = None,
        base_url: str = "https://api.mia21.com",
        user_id: Optional[str] = None,
        timeout: int = 90,
        customer_llm_key: Optional[str] = None
    ):
        """
        Initialize Mia21 client.
        
        Args:
            api_key: Your Mia21 API key (optional if using BYOK)
            base_url: API base URL (default: production)
            user_id: Unique user identifier (auto-generated if not provided)
            timeout: Request timeout in seconds (default: 90)
            customer_llm_key: Your LLM API key for BYOK (OpenAI or Gemini)
        """
        self.api_key = api_key
        self.customer_llm_key = customer_llm_key
        self.base_url = base_url.rstrip('/')
        self.api_url = f"{self.base_url}/api/v1"
        self.user_id = user_id or str(uuid.uuid4())
        self.timeout = timeout
        self.current_space = None
        self._session = requests.Session()
        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["x-api-key"] = api_key
        self._session.headers.update(headers)
    
    def list_spaces(self) -> List[Space]:
        """
        List all available spaces.
        
        Returns:
            List of Space objects
            
        Example:
            >>> spaces = client.list_spaces()
            >>> for space in spaces:
            ...     print(f"{space.id}: {space.name}")
        """
        try:
            response = self._session.post(
                f"{self.api_url}/list_spaces",
                json={"app_id": self.user_id},
                timeout=self.timeout
            )
            response.raise_for_status()
            data = response.json()
            return [Space(**s) for s in data.get("spaces", [])]
        except requests.RequestException as e:
            raise APIError(f"Failed to list spaces: {e}")
    
    def initialize(
        self,
        space_id: str = "dr_panda",
        llm_type: str = "openai",
        user_name: Optional[str] = None,
        language: Optional[str] = None,
        generate_first_message: bool = True,
        incognito_mode: bool = False,
        customer_llm_key: Optional[str] = None,
        space_config: Optional[Dict[str, Any]] = None
    ) -> InitializeResponse:
        """
        Initialize a chat session.
        
        Args:
            space_id: Space to use (default: "dr_panda")
            llm_type: "openai" or "gemini"
            user_name: User's display name
            language: Force language (e.g., "es", "de")
            generate_first_message: Generate AI greeting
            incognito_mode: Privacy mode (no data saved)
            customer_llm_key: Your LLM API key for BYOK (overrides instance key)
            space_config: Complete space configuration (for external/custom spaces)
            
        Returns:
            InitializeResponse with first message
            
        Example:
            >>> # With BYOK
            >>> response = client.initialize(
            ...     space_id="my-chatbot",
            ...     llm_type="gemini",
            ...     customer_llm_key="your-gemini-key"
            ... )
            >>> print(response.message)
            
            >>> # With custom space config
            >>> response = client.initialize(
            ...     space_id="custom_bot",
            ...     llm_type="gemini",
            ...     customer_llm_key="your-gemini-key",
            ...     space_config={
            ...         "space_id": "custom_bot",
            ...         "prompt": "You are a helpful assistant",
            ...         "description": "Custom chatbot",
            ...         "llm_identifier": "gemini-2.5-flash",
            ...         "temperature": 0.7,
            ...         "max_tokens": 1000
            ...     }
            ... )
        """
        try:
            payload = {
                "app_id": self.user_id,
                "space_id": space_id,
                "llm_type": llm_type,
                "user_name": user_name,
                "language": language,
                "generate_first_message": generate_first_message,
                "incognito_mode": incognito_mode
            }
            
            # Add customer LLM key if provided
            llm_key = customer_llm_key or self.customer_llm_key
            if llm_key:
                payload["customer_llm_key"] = llm_key
            
            # Add space config if provided (for external spaces)
            if space_config:
                payload["space_config"] = space_config
            
            response = self._session.post(
                f"{self.api_url}/initialize_chat",
                json=payload,
                timeout=self.timeout
            )
            response.raise_for_status()
            data = response.json()
            self.current_space = space_id
            return InitializeResponse(**data)
        except requests.RequestException as e:
            raise APIError(f"Failed to initialize chat: {e}")
    
    def chat(
        self,
        message: str,
        space_id: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        customer_llm_key: Optional[str] = None,
        space_config: Optional[Dict[str, Any]] = None,
        llm_type: Optional[str] = None
    ) -> ChatResponse:
        """
        Send a message and get a response.
        
        Args:
            message: User message
            space_id: Which space to chat with (uses current if not specified)
            temperature: Override temperature (0.0-2.0)
            max_tokens: Override max tokens
            customer_llm_key: Your LLM API key for BYOK (overrides instance key)
            space_config: Complete space configuration (for external spaces)
            llm_type: LLM type to use (overrides default)
            
        Returns:
            ChatResponse with AI message and tool_calls (if any)
            
        Example:
            >>> response = client.chat("I'm feeling anxious today")
            >>> print(response.message)
            >>> if response.tool_calls:
            ...     print(f"Functions triggered: {[tc['name'] for tc in response.tool_calls]}")
        """
        if not self.current_space and not space_id:
            raise ChatNotInitializedError("Chat not initialized. Call initialize() first.")
        
        try:
            payload = {
                "app_id": self.user_id,
                "space_id": space_id or self.current_space,
                "messages": [{"role": "user", "content": message}],
                "llm_type": llm_type or "openai",
                "temperature": temperature,
                "max_tokens": max_tokens,
                "stream": False
            }
            
            # Add customer LLM key if provided
            llm_key = customer_llm_key or self.customer_llm_key
            if llm_key:
                payload["customer_llm_key"] = llm_key
            
            # Add space config if provided
            if space_config:
                payload["space_config"] = space_config
            
            response = self._session.post(
                f"{self.api_url}/chat",
                json=payload,
                timeout=self.timeout
            )
            response.raise_for_status()
            data = response.json()
            return ChatResponse(**data)
        except requests.RequestException as e:
            raise APIError(f"Failed to send message: {e}")
    
    def stream_chat(
        self,
        message: str,
        space_id: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        customer_llm_key: Optional[str] = None,
        space_config: Optional[Dict[str, Any]] = None,
        llm_type: Optional[str] = None
    ) -> Generator[str, None, None]:
        """
        Send a message and stream the response in real-time.
        
        Args:
            message: User message
            space_id: Which space to chat with
            temperature: Override temperature
            max_tokens: Override max tokens
            customer_llm_key: Your LLM API key for BYOK
            space_config: Complete space configuration (for external spaces)
            llm_type: LLM type to use
            
        Yields:
            Text chunks as they arrive
            
        Example:
            >>> for chunk in client.stream_chat("Tell me a story"):
            ...     print(chunk, end='', flush=True)
            
            >>> # With custom space
            >>> for chunk in client.stream_chat(
            ...     "Hello!",
            ...     space_config={
            ...         "space_id": "my_bot",
            ...         "prompt": "You are helpful",
            ...         "llm_identifier": "gemini-2.5-flash",
            ...         "temperature": 0.7,
            ...         "max_tokens": 1000
            ...     },
            ...     llm_type="gemini",
            ...     customer_llm_key="your-gemini-key"
            ... ):
            ...     print(chunk, end='', flush=True)
        """
        if not self.current_space and not space_id:
            raise ChatNotInitializedError("Chat not initialized. Call initialize() first.")
        
        try:
            payload = {
                "app_id": self.user_id,
                "space_id": space_id or self.current_space,
                "messages": [{"role": "user", "content": message}],
                "llm_type": llm_type or "openai",
                "temperature": temperature,
                "max_tokens": max_tokens,
                "stream": True
            }
            
            # Add customer LLM key if provided
            llm_key = customer_llm_key or self.customer_llm_key
            if llm_key:
                payload["customer_llm_key"] = llm_key
            
            # Add space config if provided
            if space_config:
                payload["space_config"] = space_config
            
            response = self._session.post(
                f"{self.api_url}/chat/stream",
                json=payload,
                stream=True,
                timeout=self.timeout
            )
            response.raise_for_status()
            
            for line in response.iter_lines():
                if line:
                    line_text = line.decode('utf-8')
                    if line_text.startswith('data: '):
                        data = json.loads(line_text[6:])
                        
                        # Handle errors (skip function_call conversion errors)
                        if 'error' in data and data['error']:
                            error_msg = data['error']
                            if 'function_call' not in error_msg and 'convert' not in error_msg.lower():
                                raise APIError(f"Streaming error: {error_msg}")
                        
                        # Handle function calls (logged but not yielded)
                        if data.get('type') == 'function_call':
                            # Function calls execute silently in background
                            continue
                        
                        # Handle text content
                        if 'content' in data:
                            yield data['content']
                        
                        # Handle completion
                        if data.get('done'):
                            # Log any tool calls that were triggered
                            if data.get('tool_calls'):
                                import logging
                                logging.info(f"Functions triggered: {[tc['name'] for tc in data['tool_calls']]}")
                            break
        except requests.RequestException as e:
            raise APIError(f"Failed to stream message: {e}")
    
    def close(self, space_id: Optional[str] = None):
        """
        Close chat session and save conversation.
        
        Args:
            space_id: Which space to close (current if not specified)
            
        Example:
            >>> client.close()
        """
        try:
            response = self._session.post(
                f"{self.api_url}/close_chat",
                json={
                    "app_id": self.user_id,
                    "space_id": space_id or self.current_space
                },
                timeout=self.timeout
            )
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            raise APIError(f"Failed to close chat: {e}")
    
    def __enter__(self):
        """Context manager support."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Auto-close on context exit."""
        if self.current_space:
            try:
                self.close()
            except:
                pass

