import os
import asyncio
import requests
from typing import List, Dict, Any, Iterator, AsyncIterator
from litellm.types.utils import GenericStreamingChunk, ModelResponse
from litellm.llms.custom_llm import CustomLLM
from litellm import completion, acompletion


class HumanAdapter(CustomLLM):
    """Human Adapter for human-in-the-loop completions.
    
    This adapter sends completion requests to a human completion server
    where humans can review and respond to AI requests.
    """
    
    def __init__(self, base_url: str | None = None, timeout: float = 300.0, **kwargs):
        """Initialize the human adapter.
        
        Args:
            base_url: Base URL for the human completion server.
                     Defaults to HUMAN_BASE_URL environment variable or http://localhost:8002
            timeout: Timeout in seconds for waiting for human response
            **kwargs: Additional arguments
        """
        super().__init__()
        self.base_url = base_url or os.getenv('HUMAN_BASE_URL', 'http://localhost:8002')
        self.timeout = timeout
        
        # Ensure base_url doesn't end with slash
        self.base_url = self.base_url.rstrip('/')
    
    def _queue_completion(self, messages: List[Dict[str, Any]], model: str) -> str:
        """Queue a completion request and return the call ID.
        
        Args:
            messages: Messages in OpenAI format
            model: Model name
            
        Returns:
            Call ID for tracking the request
            
        Raises:
            Exception: If queueing fails
        """
        try:
            response = requests.post(
                f"{self.base_url}/queue",
                json={"messages": messages, "model": model},
                timeout=10
            )
            response.raise_for_status()
            return response.json()["id"]
        except requests.RequestException as e:
            raise Exception(f"Failed to queue completion request: {e}")
    
    def _wait_for_completion(self, call_id: str) -> Dict[str, Any]:
        """Wait for human to complete the call.
        
        Args:
            call_id: ID of the queued completion call
            
        Returns:
            Dict containing response and/or tool_calls
            
        Raises:
            TimeoutError: If timeout is exceeded
            Exception: If completion fails
        """
        import time
        
        start_time = time.time()
        
        while True:
            try:
                # Check status
                status_response = requests.get(f"{self.base_url}/status/{call_id}")
                status_response.raise_for_status()
                status_data = status_response.json()
                
                if status_data["status"] == "completed":
                    result = {}
                    if "response" in status_data and status_data["response"]:
                        result["response"] = status_data["response"]
                    if "tool_calls" in status_data and status_data["tool_calls"]:
                        result["tool_calls"] = status_data["tool_calls"]
                    return result
                elif status_data["status"] == "failed":
                    error_msg = status_data.get("error", "Unknown error")
                    raise Exception(f"Completion failed: {error_msg}")
                
                # Check timeout
                if time.time() - start_time > self.timeout:
                    raise TimeoutError(f"Timeout waiting for human response after {self.timeout} seconds")
                
                # Wait before checking again
                time.sleep(1.0)
                
            except requests.RequestException as e:
                if time.time() - start_time > self.timeout:
                    raise TimeoutError(f"Timeout waiting for human response: {e}")
                # Continue trying if we haven't timed out
                time.sleep(1.0)
    
    async def _async_wait_for_completion(self, call_id: str) -> Dict[str, Any]:
        """Async version of wait_for_completion.
        
        Args:
            call_id: ID of the queued completion call
            
        Returns:
            Dict containing response and/or tool_calls
            
        Raises:
            TimeoutError: If timeout is exceeded
            Exception: If completion fails
        """
        import aiohttp
        import time
        
        start_time = time.time()
        
        async with aiohttp.ClientSession() as session:
            while True:
                try:
                    # Check status
                    async with session.get(f"{self.base_url}/status/{call_id}") as response:
                        response.raise_for_status()
                        status_data = await response.json()
                    
                    if status_data["status"] == "completed":
                        result = {}
                        if "response" in status_data and status_data["response"]:
                            result["response"] = status_data["response"]
                        if "tool_calls" in status_data and status_data["tool_calls"]:
                            result["tool_calls"] = status_data["tool_calls"]
                        return result
                    elif status_data["status"] == "failed":
                        error_msg = status_data.get("error", "Unknown error")
                        raise Exception(f"Completion failed: {error_msg}")
                    
                    # Check timeout
                    if time.time() - start_time > self.timeout:
                        raise TimeoutError(f"Timeout waiting for human response after {self.timeout} seconds")
                    
                    # Wait before checking again
                    await asyncio.sleep(1.0)
                    
                except Exception as e:
                    if time.time() - start_time > self.timeout:
                        raise TimeoutError(f"Timeout waiting for human response: {e}")
                    # Continue trying if we haven't timed out
                    await asyncio.sleep(1.0)
    
    def _generate_response(self, messages: List[Dict[str, Any]], model: str) -> Dict[str, Any]:
        """Generate a human response for the given messages.
        
        Args:
            messages: Messages in OpenAI format
            model: Model name
            
        Returns:
            Dict containing response and/or tool_calls
        """
        # Queue the completion request
        call_id = self._queue_completion(messages, model)
        
        # Wait for human response
        response = self._wait_for_completion(call_id)
        
        return response
    
    async def _async_generate_response(self, messages: List[Dict[str, Any]], model: str) -> Dict[str, Any]:
        """Async version of _generate_response.
        
        Args:
            messages: Messages in OpenAI format
            model: Model name
            
        Returns:
            Dict containing response and/or tool_calls
        """
        # Queue the completion request (sync operation)
        call_id = self._queue_completion(messages, model)
        
        # Wait for human response (async)
        response = await self._async_wait_for_completion(call_id)
        
        return response
    
    def completion(self, *args, **kwargs) -> ModelResponse:
        """Synchronous completion method.
        
        Returns:
            ModelResponse with human-generated text or tool calls
        """
        messages = kwargs.get('messages', [])
        model = kwargs.get('model', 'human')
        
        # Generate human response
        human_response_data = self._generate_response(messages, model)
        
        # Create ModelResponse with proper structure
        from litellm.types.utils import ModelResponse, Choices, Message
        import uuid
        import time
        
        # Create message content based on response type
        if "tool_calls" in human_response_data and human_response_data["tool_calls"]:
            # Tool calls response
            message = Message(
                role="assistant",
                content=human_response_data.get("response", ""),
                tool_calls=human_response_data["tool_calls"]
            )
        else:
            # Text response
            message = Message(
                role="assistant",
                content=human_response_data.get("response", "")
            )
        
        choice = Choices(
            finish_reason="stop",
            index=0,
            message=message
        )
        
        result = ModelResponse(
            id=f"human-{uuid.uuid4()}",
            choices=[choice],
            created=int(time.time()),
            model=f"human/{model}",
            object="chat.completion"
        )
        
        return result
    
    async def acompletion(self, *args, **kwargs) -> ModelResponse:
        """Asynchronous completion method.
        
        Returns:
            ModelResponse with human-generated text or tool calls
        """
        messages = kwargs.get('messages', [])
        model = kwargs.get('model', 'human')
        
        # Generate human response
        human_response_data = await self._async_generate_response(messages, model)
        
        # Create ModelResponse with proper structure
        from litellm.types.utils import ModelResponse, Choices, Message
        import uuid
        import time
        
        # Create message content based on response type
        if "tool_calls" in human_response_data and human_response_data["tool_calls"]:
            # Tool calls response
            message = Message(
                role="assistant",
                content=human_response_data.get("response", ""),
                tool_calls=human_response_data["tool_calls"]
            )
        else:
            # Text response
            message = Message(
                role="assistant",
                content=human_response_data.get("response", "")
            )
        
        choice = Choices(
            finish_reason="stop",
            index=0,
            message=message
        )
        
        result = ModelResponse(
            id=f"human-{uuid.uuid4()}",
            choices=[choice],
            created=int(time.time()),
            model=f"human/{model}",
            object="chat.completion"
        )
        
        return result
    
    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
        """Synchronous streaming method.
        
        Yields:
            Streaming chunks with human-generated text or tool calls
        """
        messages = kwargs.get('messages', [])
        model = kwargs.get('model', 'human')
        
        # Generate human response
        human_response_data = self._generate_response(messages, model)
        
        import time
        
        # Handle tool calls vs text response
        if "tool_calls" in human_response_data and human_response_data["tool_calls"]:
            # Stream tool calls as a single chunk
            generic_chunk: GenericStreamingChunk = {
                "finish_reason": "tool_calls",
                "index": 0,
                "is_finished": True,
                "text": human_response_data.get("response", ""),
                "tool_use": human_response_data["tool_calls"],
                "usage": {"completion_tokens": 1, "prompt_tokens": 0, "total_tokens": 1},
            }
            yield generic_chunk
        else:
            # Stream text response
            response_text = human_response_data.get("response", "")
            generic_chunk: GenericStreamingChunk = {
                "finish_reason": "stop",
                "index": 0,
                "is_finished": True,
                "text": response_text,
                "tool_use": None,
                "usage": {"completion_tokens": len(response_text.split()), "prompt_tokens": 0, "total_tokens": len(response_text.split())},
            }
            yield generic_chunk
    
    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
        """Asynchronous streaming method.
        
        Yields:
            Streaming chunks with human-generated text or tool calls
        """
        messages = kwargs.get('messages', [])
        model = kwargs.get('model', 'human')
        
        # Generate human response
        human_response = await self._async_generate_response(messages, model)
        
        # Return as single streaming chunk
        generic_streaming_chunk: GenericStreamingChunk = {
            "finish_reason": "stop",
            "index": 0,
            "is_finished": True,
            "text": human_response,
            "tool_use": None,
            "usage": {"completion_tokens": len(human_response.split()), "prompt_tokens": 0, "total_tokens": len(human_response.split())},
        }
        
        yield generic_streaming_chunk