"""Simons AI Assistant client for the CPZ SDK.

Provides synchronous and streaming access to the Simons AI Assistant,
with full support for chat, memory management, and code proposals.

Example:
    from cpz import CPZClient
    
    client = CPZClient()
    
    # Simple chat
    response = client.simons.chat("Analyze AAPL for momentum trading")
    print(response.content)
    
    # Streaming response
    for chunk in client.simons.stream("Write a backtest for RSI strategy"):
        print(chunk.content, end="", flush=True)
"""
from __future__ import annotations

import json
from typing import Any, Dict, Iterator, Literal, Optional

import requests

from .common.cpz_ai import CPZAIClient
from .common.logging import get_logger
from .simons_models import SimonsChunk, SimonsMemory, SimonsResponse


class SimonsClient:
    """Client for Simons AI Assistant.
    
    Provides methods to interact with the Simons AI through the CPZ SDK,
    including synchronous chat, streaming responses, and memory management.
    
    All requests are authenticated using your CPZ API key/secret and
    AI usage is tracked against your account's spending limits.
    """
    
    # Default endpoint (can be overridden for testing)
    DEFAULT_ENDPOINT = "https://api-ai.cpz-lab.com/cpz/simons"
    
    def __init__(self, cpz_client: CPZAIClient) -> None:
        """Initialize the Simons client.
        
        Args:
            cpz_client: The CPZAIClient instance for authentication
        """
        self._client = cpz_client
        self._logger = get_logger()
        self._endpoint = self.DEFAULT_ENDPOINT
    
    def _get_headers(self) -> Dict[str, str]:
        """Get headers for API requests."""
        return {
            "X-CPZ-Key": self._client.api_key,
            "X-CPZ-Secret": self._client.secret_key,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
    
    def _get_stream_headers(self) -> Dict[str, str]:
        """Get headers for streaming API requests."""
        headers = self._get_headers()
        headers["Accept"] = "text/event-stream"
        return headers
    
    def chat(
        self,
        message: str,
        *,
        mode: Literal["ask", "agent", "debug"] = "ask",
        strategy_id: Optional[str] = None,
        model: Optional[str] = None,
        show_thinking: bool = True,
        timeout: float = 120.0,
    ) -> SimonsResponse:
        """Send a message to Simons and get a complete response.
        
        Args:
            message: The message/question to send to Simons
            mode: The assistant mode:
                - "ask": Read-only assistance (questions, analysis)
                - "agent": Can propose and apply code changes
                - "debug": Debugging assistance with execution context
            strategy_id: Optional strategy ID for context (allows file access)
            model: AI model to use (default: claude-sonnet-4-5-20250929)
            show_thinking: Whether to include extended thinking in response
            timeout: Request timeout in seconds (default: 120s for AI responses)
        
        Returns:
            SimonsResponse with the AI's response, usage stats, and any proposals
        
        Example:
            response = client.simons.chat("Explain moving average crossover strategy")
            print(response.content)
            print(f"Cost: ${response.cost:.4f}")
        """
        payload = {
            "message": message,
            "mode": mode,
            "show_thinking": show_thinking,
        }
        
        if strategy_id:
            payload["strategy_id"] = strategy_id
        if model:
            payload["model"] = model
        
        try:
            response = requests.post(
                f"{self._endpoint}/chat",
                headers=self._get_headers(),
                json=payload,
                timeout=timeout,
            )
            
            if response.status_code == 403:
                error_data = response.json()
                raise PermissionError(
                    error_data.get("message", "Simons access not authorized. "
                    "Enable the 'simons' scope in Settings > API Keys.")
                )
            
            response.raise_for_status()
            data = response.json()
            
            return SimonsResponse.from_dict(data)
            
        except requests.exceptions.Timeout:
            self._logger.error("simons_chat_timeout", timeout=timeout)
            raise TimeoutError(f"Simons request timed out after {timeout}s")
        except requests.exceptions.RequestException as e:
            self._logger.error("simons_chat_error", error=str(e))
            raise
    
    def stream(
        self,
        message: str,
        *,
        mode: Literal["ask", "agent", "debug"] = "ask",
        strategy_id: Optional[str] = None,
        model: Optional[str] = None,
        show_thinking: bool = True,
        timeout: float = 120.0,
    ) -> Iterator[SimonsChunk]:
        """Send a message to Simons and stream the response.
        
        Yields chunks as they arrive, allowing real-time display of the response.
        This is similar to how Claude's streaming API works.
        
        Args:
            message: The message/question to send to Simons
            mode: The assistant mode ("ask", "agent", or "debug")
            strategy_id: Optional strategy ID for context
            model: AI model to use
            show_thinking: Whether to include extended thinking chunks
            timeout: Request timeout in seconds
        
        Yields:
            SimonsChunk objects with type ('text', 'thinking', 'done') and content
        
        Example:
            for chunk in client.simons.stream("Write a backtest"):
                if chunk.type == "text":
                    print(chunk.content, end="", flush=True)
                elif chunk.type == "done":
                    print()  # Newline at end
        """
        payload = {
            "message": message,
            "mode": mode,
            "show_thinking": show_thinking,
        }
        
        if strategy_id:
            payload["strategy_id"] = strategy_id
        if model:
            payload["model"] = model
        
        try:
            response = requests.post(
                f"{self._endpoint}/stream",
                headers=self._get_stream_headers(),
                json=payload,
                timeout=timeout,
                stream=True,
            )
            
            if response.status_code == 403:
                try:
                    error_data = response.json()
                    raise PermissionError(
                        error_data.get("message", "Simons access not authorized.")
                    )
                except json.JSONDecodeError:
                    raise PermissionError("Simons access not authorized.")
            
            response.raise_for_status()
            
            # Parse Server-Sent Events
            event_type = "text"
            for line in response.iter_lines(decode_unicode=True):
                if not line:
                    continue
                
                if line.startswith("event:"):
                    event_type = line[6:].strip()
                elif line.startswith("data:"):
                    data = line[5:].strip()
                    
                    if data == "[DONE]":
                        yield SimonsChunk(type="done", content="", delta="")
                        break
                    
                    # Try to parse as JSON for structured chunks
                    try:
                        parsed = json.loads(data)
                        # Support multiple field names: token, content, delta, text
                        content = (
                            parsed.get("token") or 
                            parsed.get("content") or 
                            parsed.get("delta") or 
                            parsed.get("text") or 
                            ""
                        )
                        # Determine chunk type from response
                        chunk_type = parsed.get("type", event_type)
                        # Check for thinking content
                        if parsed.get("thinking"):
                            chunk_type = "thinking"
                            content = parsed.get("thinking")
                        
                        if content:  # Only yield if there's actual content
                            yield SimonsChunk(
                                type=chunk_type if chunk_type in ("text", "thinking", "done", "error") else "text",
                                content=content,
                                delta=content,
                            )
                    except json.JSONDecodeError:
                        # Plain text chunk
                        if data.strip():
                            yield SimonsChunk(type="text", content=data, delta=data)
                        
        except requests.exceptions.Timeout:
            self._logger.error("simons_stream_timeout", timeout=timeout)
            yield SimonsChunk(type="error", content=f"Request timed out after {timeout}s", delta="")
        except requests.exceptions.RequestException as e:
            self._logger.error("simons_stream_error", error=str(e))
            yield SimonsChunk(type="error", content=str(e), delta="")
    
    def get_memory(self, timeout: float = 10.0) -> SimonsMemory:
        """Get your Simons memory (preferences, learned facts).
        
        Args:
            timeout: Request timeout in seconds
        
        Returns:
            SimonsMemory object with your preferences and facts
        
        Example:
            memory = client.simons.get_memory()
            print(f"Risk profile: {memory.preferences.get('risk_profile')}")
            print(f"Learned facts: {memory.facts}")
        """
        try:
            response = requests.get(
                f"{self._endpoint}/memory",
                headers=self._get_headers(),
                timeout=timeout,
            )
            
            response.raise_for_status()
            data = response.json()
            
            return SimonsMemory.from_dict(data)
            
        except requests.exceptions.RequestException as e:
            self._logger.warning("simons_get_memory_error", error=str(e))
            return SimonsMemory()
    
    def update_memory(
        self,
        preferences: Optional[Dict[str, Any]] = None,
        facts: Optional[list[str]] = None,
        timeout: float = 10.0,
    ) -> bool:
        """Update your Simons memory.
        
        Args:
            preferences: Dictionary of preferences to update/merge
            facts: List of facts to set (replaces existing facts)
            timeout: Request timeout in seconds
        
        Returns:
            True if update was successful
        
        Example:
            client.simons.update_memory(
                preferences={"risk_profile": "moderate", "preferred_sectors": ["tech"]},
                facts=["Prefers momentum strategies", "Uses 1% position sizing"]
            )
        """
        payload: Dict[str, Any] = {}
        
        if preferences is not None:
            payload["preferences"] = preferences
        if facts is not None:
            payload["facts"] = facts
        
        if not payload:
            return True  # Nothing to update
        
        try:
            response = requests.post(
                f"{self._endpoint}/memory",
                headers=self._get_headers(),
                json=payload,
                timeout=timeout,
            )
            
            response.raise_for_status()
            data = response.json()
            
            return data.get("success", True)
            
        except requests.exceptions.RequestException as e:
            self._logger.error("simons_update_memory_error", error=str(e))
            return False


class AsyncSimonsClient:
    """Async client for Simons AI Assistant.
    
    Provides async methods to interact with the Simons AI.
    Use this with asyncio for non-blocking operations.
    """
    
    DEFAULT_ENDPOINT = "https://api-ai.cpz-lab.com/cpz/simons"
    
    def __init__(self, cpz_client: CPZAIClient) -> None:
        """Initialize the async Simons client.
        
        Args:
            cpz_client: The CPZAIClient instance for authentication
        """
        self._client = cpz_client
        self._logger = get_logger()
        self._endpoint = self.DEFAULT_ENDPOINT
    
    def _get_headers(self) -> Dict[str, str]:
        """Get headers for API requests."""
        return {
            "X-CPZ-Key": self._client.api_key,
            "X-CPZ-Secret": self._client.secret_key,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
    
    async def chat(
        self,
        message: str,
        *,
        mode: Literal["ask", "agent", "debug"] = "ask",
        strategy_id: Optional[str] = None,
        model: Optional[str] = None,
        show_thinking: bool = True,
        timeout: float = 120.0,
    ) -> SimonsResponse:
        """Send a message to Simons and get a complete response (async).
        
        See SimonsClient.chat for full documentation.
        """
        import httpx
        
        payload = {
            "message": message,
            "mode": mode,
            "show_thinking": show_thinking,
        }
        
        if strategy_id:
            payload["strategy_id"] = strategy_id
        if model:
            payload["model"] = model
        
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(
                f"{self._endpoint}/chat",
                headers=self._get_headers(),
                json=payload,
            )
            
            if response.status_code == 403:
                error_data = response.json()
                raise PermissionError(
                    error_data.get("message", "Simons access not authorized.")
                )
            
            response.raise_for_status()
            data = response.json()
            
            return SimonsResponse.from_dict(data)
    
    async def stream(
        self,
        message: str,
        *,
        mode: Literal["ask", "agent", "debug"] = "ask",
        strategy_id: Optional[str] = None,
        model: Optional[str] = None,
        show_thinking: bool = True,
        timeout: float = 120.0,
    ):
        """Send a message to Simons and stream the response (async).
        
        See SimonsClient.stream for full documentation.
        
        Yields:
            SimonsChunk objects with type and content
        """
        import httpx
        
        payload = {
            "message": message,
            "mode": mode,
            "show_thinking": show_thinking,
        }
        
        if strategy_id:
            payload["strategy_id"] = strategy_id
        if model:
            payload["model"] = model
        
        headers = self._get_headers()
        headers["Accept"] = "text/event-stream"
        
        async with httpx.AsyncClient(timeout=timeout) as client:
            async with client.stream(
                "POST",
                f"{self._endpoint}/stream",
                headers=headers,
                json=payload,
            ) as response:
                if response.status_code == 403:
                    raise PermissionError("Simons access not authorized.")
                
                response.raise_for_status()
                
                event_type = "text"
                async for line in response.aiter_lines():
                    if not line:
                        continue
                    
                    if line.startswith("event:"):
                        event_type = line[6:].strip()
                    elif line.startswith("data:"):
                        data = line[5:].strip()
                        
                        if data == "[DONE]":
                            yield SimonsChunk(type="done", content="", delta="")
                            break
                        
                        try:
                            parsed = json.loads(data)
                            # Support multiple field names: token, content, delta, text
                            content = (
                                parsed.get("token") or 
                                parsed.get("content") or 
                                parsed.get("delta") or 
                                parsed.get("text") or 
                                ""
                            )
                            chunk_type = parsed.get("type", event_type)
                            # Check for thinking content
                            if parsed.get("thinking"):
                                chunk_type = "thinking"
                                content = parsed.get("thinking")
                            
                            if content:  # Only yield if there's actual content
                                yield SimonsChunk(
                                    type=chunk_type if chunk_type in ("text", "thinking", "done", "error") else "text",
                                    content=content,
                                    delta=content,
                                )
                        except json.JSONDecodeError:
                            if data.strip():
                                yield SimonsChunk(type="text", content=data, delta=data)
    
    async def get_memory(self, timeout: float = 10.0) -> SimonsMemory:
        """Get your Simons memory (async).
        
        See SimonsClient.get_memory for full documentation.
        """
        import httpx
        
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.get(
                f"{self._endpoint}/memory",
                headers=self._get_headers(),
            )
            
            response.raise_for_status()
            data = response.json()
            
            return SimonsMemory.from_dict(data)
    
    async def update_memory(
        self,
        preferences: Optional[Dict[str, Any]] = None,
        facts: Optional[list[str]] = None,
        timeout: float = 10.0,
    ) -> bool:
        """Update your Simons memory (async).
        
        See SimonsClient.update_memory for full documentation.
        """
        import httpx
        
        payload: Dict[str, Any] = {}
        
        if preferences is not None:
            payload["preferences"] = preferences
        if facts is not None:
            payload["facts"] = facts
        
        if not payload:
            return True
        
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(
                f"{self._endpoint}/memory",
                headers=self._get_headers(),
                json=payload,
            )
            
            response.raise_for_status()
            data = response.json()
            
            return data.get("success", True)
