"""OpenAI Response API implementation"""

import json
import logging
import time
from typing import Optional, Type, Dict, Any
from pydantic import BaseModel
from ..models import TextRequest, MessageRequest, TextResponse, StreamChunk
from ..utils import JSONFileCache

logger = logging.getLogger('aws_llm_wrapper')


class ResponsesAPI:
    """Handler for OpenAI Response API"""
    
    def __init__(self, client, config, cache: JSONFileCache, semaphore=None):
        self.client = client
        self.config = config
        self.cache = cache
        self.semaphore = semaphore
    
    async def generate_text(self, request: TextRequest, invoke_with_retry) -> TextResponse:
        """Generate text using Response API"""
        model = request.model or self.config.default_model
        is_reasoning = request.reasoning_effort is not None
        
        # Reasoning models don't support temperature
        if is_reasoning and request.temperature is not None and request.temperature != 1:
            raise ValueError(
                f"Reasoning models do not support temperature. "
                f"Remove temperature from your request or set it to 1."
            )
        
        temperature = None if is_reasoning else (request.temperature if request.temperature is not None else 0)
        
        # Cache key - reasoning models always cache (no temperature variation)
        cache_key = None
        if not request.stream and (temperature == 0 or is_reasoning):
            cache_key = self.cache._generate_key(
                api_type="responses",
                model=model,
                prompt=request.prompt,
                max_tokens=request.max_tokens or self.config.max_tokens,
                instructions=request.system_prompt,
                reasoning_effort=request.reasoning_effort,
                response_format=request.response_format.__name__ if request.response_format else None
            )
        
        if request.clear_cache and cache_key:
            self.cache.clear(cache_key)
            logger.info(f"Cleared cache entry: {cache_key[:8]}...")
        
        if request.use_cache and cache_key:
            cached = self.cache.get(cache_key)
            if cached:
                logger.info(f"Cache hit [{cache_key[:8]}] - {model} - prompt: {request.prompt[:50]}...")
                return self._deserialize_response(cached["data"], request.response_format)
        
        prompt_preview = request.prompt[:60] + "..." if len(request.prompt) > 60 else request.prompt
        logger.info(f"API call to {model} (Response API) - reasoning={request.reasoning_effort or 'off'} - prompt: {prompt_preview}")
        
        start_time = time.time()
        
        # Build params
        params = {
            "model": model,
            "input": request.prompt,
        }
        
        if temperature is not None:
            params["temperature"] = temperature
        
        if request.reasoning_effort:
            params["reasoning"] = {"effort": request.reasoning_effort}
        
        if request.system_prompt:
            params["instructions"] = request.system_prompt
        
        if request.max_tokens:
            params["max_output_tokens"] = request.max_tokens
        
        if request.top_p:
            params["top_p"] = request.top_p
        
        # Structured output
        if request.response_format:
            schema = request.response_format.model_json_schema()
            schema["additionalProperties"] = False
            params["text"] = {
                "format": {
                    "type": "json_schema",
                    "name": request.response_format.__name__,
                    "schema": schema,
                    "strict": True
                }
            }
        
        try:
            if self.semaphore:
                async with self.semaphore:
                    response = await invoke_with_retry(self.client.responses.create, **params)
            else:
                response = await invoke_with_retry(self.client.responses.create, **params)
            
            result = self._parse_response(response, model, request.response_format)
            
            elapsed = time.time() - start_time
            logger.info(
                f"Response received - {result.input_tokens} in / {result.output_tokens} out tokens - "
                f"{elapsed:.2f}s - {result.text[:50]}..."
            )
            
            if cache_key:
                self.cache.set(cache_key, self._serialize_response(result), {})
                logger.debug(f"Cached response: {cache_key[:8]}...")
            
            return result
        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(f"Error after {elapsed:.2f}s - {model}: {str(e)}")
            raise
    
    def _parse_response(self, response, model: str, response_format: Optional[Type[BaseModel]] = None) -> TextResponse:
        """Parse Response API response"""
        text = response.output_text or ""
        structured_data = None
        
        if response_format and text:
            try:
                data = json.loads(text)
                structured_data = response_format(**data)
            except Exception as e:
                logger.warning(f"Failed to parse structured output: {e}")
        
        input_tokens = response.usage.input_tokens if response.usage else 0
        output_tokens = response.usage.output_tokens if response.usage else 0
        
        logger.debug(f"Raw usage: {response.usage}")
        
        # Capture reasoning tokens in metadata if present
        metadata = {}
        if response.usage:
            reasoning_tokens = getattr(getattr(response.usage, "output_tokens_details", None), "reasoning_tokens", 0)
            cached_tokens = getattr(getattr(response.usage, "input_tokens_details", None), "cached_tokens", 0)
            if reasoning_tokens:
                metadata["reasoning_tokens"] = reasoning_tokens
            if cached_tokens:
                metadata["cached_tokens"] = cached_tokens
        
        return TextResponse(
            text=text,
            model=model,
            stop_reason=response.status or "",
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            metadata=metadata,
            structured_data=structured_data,
        )
    
    def _serialize_response(self, response: TextResponse) -> Dict[str, Any]:
        """Serialize TextResponse for caching"""
        return {
            "text": response.text,
            "model": response.model,
            "stop_reason": response.stop_reason,
            "input_tokens": response.input_tokens,
            "output_tokens": response.output_tokens,
            "metadata": response.metadata,
            "structured_data": response.structured_data.model_dump() if response.structured_data else None,
        }
    
    def _deserialize_response(self, data: Dict[str, Any], response_format: Optional[Type[BaseModel]] = None) -> TextResponse:
        """Deserialize cached data back to TextResponse"""
        structured_data = None
        if data.get("structured_data") and response_format:
            structured_data = response_format(**data["structured_data"])
        
        return TextResponse(
            text=data["text"],
            model=data["model"],
            stop_reason=data["stop_reason"],
            input_tokens=data["input_tokens"],
            output_tokens=data["output_tokens"],
            metadata=data.get("metadata", {}),
            structured_data=structured_data,
        )
