"""
Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>

AISBF - AI Service Broker Framework || AI Should Be Free

OpenAI provider handler.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.

Why did the programmer quit his job? Because he didn't get arrays!
"""
import time
from typing import Dict, List, Optional, Union
from openai import OpenAI
from ..models import Model
from ..config import config
from ..utils import count_messages_tokens
from .base import (
    BaseProviderHandler,
    AISBF_DEBUG,
    adapt_request_for_unsupported_param,
    apply_max_tokens_param,
    apply_temperature_param,
)


class OpenAIProviderHandler(BaseProviderHandler):
    def __init__(self, provider_id: str, api_key: str, user_id: Optional[int] = None, provider_config=None):
        self.provider_config = provider_config if provider_config is not None else config.providers[provider_id]
        super().__init__(provider_id, api_key, user_id=user_id)
        endpoint = self.provider_config.get("endpoint") if isinstance(self.provider_config, dict) else self.provider_config.endpoint
        self.client = OpenAI(base_url=endpoint, api_key=api_key)
    
    def validate_credentials(self) -> bool:
        """Validate OpenAI API key presence."""
        if not self.api_key:
            logging.error(f"[{self.provider_id}] API key required but not provided")
            return False
        
        stripped = self.api_key.strip()
        if not stripped or stripped.startswith('YOUR_'):
            logging.error(f"[{self.provider_id}] API key appears to be a placeholder")
            return False
        
        logging.info(f"[{self.provider_id}] API key validated")
        return True

    async def handle_request(self, model: str, messages: List[Dict], max_tokens: Optional[int] = None,
                           temperature: Optional[float] = 1.0, stream: Optional[bool] = False,
                           tools: Optional[List[Dict]] = None, tool_choice: Optional[Union[str, Dict]] = None) -> Union[Dict, object]:
        if self.is_rate_limited():
            raise Exception("Provider rate limited")

        try:
            import logging
            logging.info(f"OpenAIProviderHandler: Handling request for model {model}")
            if AISBF_DEBUG:
                logging.info(f"OpenAIProviderHandler: Messages: {messages}")
            else:
                logging.info(f"OpenAIProviderHandler: Messages count: {len(messages)}")
            if AISBF_DEBUG:
                logging.info(f"OpenAIProviderHandler: Tools: {tools}")
                logging.info(f"OpenAIProviderHandler: Tool choice: {tool_choice}")

            # Apply rate limiting
            await self.apply_rate_limit()

            # Check if native caching is enabled for this provider
            provider_config = config.providers.get(self.provider_id)
            enable_native_caching = getattr(provider_config, 'enable_native_caching', False)
            min_cacheable_tokens = getattr(provider_config, 'min_cacheable_tokens', 1024)
            prompt_cache_key = getattr(provider_config, 'prompt_cache_key', None)
            
            # Check user's cache settings (overrides provider config)
            if enable_native_caching and hasattr(self, 'user_id'):
                if not self._native_cache_user_allows(model):
                    enable_native_caching = False
                    logging.info(f"User {self.user_id} disabled cache for provider={self.provider_id}, model={model}")

            logging.info(f"OpenAIProviderHandler: Native caching enabled: {enable_native_caching}")
            if enable_native_caching:
                logging.info(f"OpenAIProviderHandler: Min cacheable tokens: {min_cacheable_tokens}, prompt_cache_key: {prompt_cache_key}")

            # Build request parameters
            request_params = {
                "model": model,
                "messages": [],
                "stream": stream
            }

            # Add temperature and the output-token limit only in the form this
            # model accepts (see the helpers for why they differ per model).
            apply_temperature_param(request_params, model, temperature)
            apply_max_tokens_param(request_params, model, max_tokens)

            # Add prompt_cache_key if provided (for OpenAI's load balancer routing optimization)
            if enable_native_caching and prompt_cache_key:
                request_params["prompt_cache_key"] = prompt_cache_key
                logging.info(f"OpenAIProviderHandler: Added prompt_cache_key to request")
            
            # Build messages with all fields (including tool_calls, tool_call_id, and cache_control)
            if enable_native_caching:
                # Count cumulative tokens for cache decision
                cumulative_tokens = 0
                for i, msg in enumerate(messages):
                    # Count tokens in this message
                    message_tokens = count_messages_tokens([msg], model)
                    cumulative_tokens += message_tokens

                    message = {"role": msg["role"]}
                    
                    # For tool role, tool_call_id is required
                    if msg["role"] == "tool":
                        if "tool_call_id" in msg and msg["tool_call_id"] is not None:
                            message["tool_call_id"] = msg["tool_call_id"]
                        else:
                            # Skip tool messages without tool_call_id
                            logging.warning(f"Skipping tool message without tool_call_id: {msg}")
                            continue
                    
                    if "content" in msg and msg["content"] is not None:
                        message["content"] = msg["content"]
                    if "tool_calls" in msg and msg["tool_calls"] is not None:
                        message["tool_calls"] = msg["tool_calls"]
                    if "name" in msg and msg["name"] is not None:
                        message["name"] = msg["name"]
                    
                    # Apply cache_control based on position and token count
                    if (msg["role"] == "system" or
                        (i < len(messages) - 2 and cumulative_tokens >= min_cacheable_tokens)):
                        message["cache_control"] = {"type": "ephemeral"}
                        logging.info(f"OpenAIProviderHandler: Applied cache_control to message {i} ({message_tokens} tokens, cumulative: {cumulative_tokens})")
                    else:
                        logging.info(f"OpenAIProviderHandler: Not caching message {i} ({message_tokens} tokens, cumulative: {cumulative_tokens})")
                    
                    request_params["messages"].append(message)
            else:
                # Standard message formatting without caching
                for msg in messages:
                    message = {"role": msg["role"]}
                    
                    # For tool role, tool_call_id is required
                    if msg["role"] == "tool":
                        if "tool_call_id" in msg and msg["tool_call_id"] is not None:
                            message["tool_call_id"] = msg["tool_call_id"]
                        else:
                            # Skip tool messages without tool_call_id
                            logging.warning(f"Skipping tool message without tool_call_id: {msg}")
                            continue
                    
                    if "content" in msg and msg["content"] is not None:
                        message["content"] = msg["content"]
                    if "tool_calls" in msg and msg["tool_calls"] is not None:
                        message["tool_calls"] = msg["tool_calls"]
                    if "name" in msg and msg["name"] is not None:
                        message["name"] = msg["name"]
                    request_params["messages"].append(message)
            
            # Add tools and tool_choice if provided
            if tools is not None:
                request_params["tools"] = tools
            if tool_choice is not None:
                request_params["tool_choice"] = tool_choice

            # Retry once per adaptable parameter (max_tokens, temperature) when
            # the upstream rejects it; adapt_request_for_unsupported_param
            # returns None for any other error, which we re-raise.
            for _ in range(2):
                try:
                    response = self.client.chat.completions.create(**request_params)
                    break
                except Exception as e:
                    adaptation = adapt_request_for_unsupported_param(request_params, model, e)
                    if adaptation is None:
                        raise
                    logging.info(f"OpenAIProviderHandler: {model} rejected a parameter, retrying ({adaptation})")
            else:
                response = self.client.chat.completions.create(**request_params)
            logging.info(f"OpenAIProviderHandler: Response received: {response}")
            # Streaming returns a lazy iterator; the upstream call has not been
            # consumed yet, so recording success here would prematurely reset the
            # failure counter. The caller records success after priming/consuming.
            if not stream:
                self.record_success()
            
            # Dump raw response if AISBF_DEBUG is enabled
            if AISBF_DEBUG:
                logging.info(f"=== RAW OPENAI RESPONSE ===")
                logging.info(f"Raw response type: {type(response)}")
                logging.info(f"Raw response: {response}")
                logging.info(f"=== END RAW OPENAI RESPONSE ===")
            
            # Return raw response without any parsing or modification
            logging.info(f"OpenAIProviderHandler: Returning raw response without parsing")
            return response
        except Exception as e:
            import logging
            logging.error(f"OpenAIProviderHandler: Error: {str(e)}", exc_info=True)
            self.record_failure()
            raise e

    async def get_models(self) -> List[Model]:
        try:
            import logging
            logging.info("OpenAIProviderHandler: Getting models list")

            # Apply rate limiting
            await self.apply_rate_limit()

            models = self.client.models.list()
            if AISBF_DEBUG:
                response_str = str(models)
                if len(response_str) > 1024:
                    response_str = response_str[:1024] + f" ... [TRUNCATED, total length: {len(response_str)} chars]"
                logging.info(f"OpenAIProviderHandler: Models received: {response_str}")
            else:
                model_count = len(models.data) if hasattr(models, 'data') else len(models) if isinstance(models, (list, dict)) else 'N/A'
                logging.info(f"OpenAIProviderHandler: Models received: {model_count} models")

            result = []
            for model in models:
                # Extract context size if available - check multiple field names
                context_size = None
                if hasattr(model, 'context_window') and model.context_window:
                    context_size = model.context_window
                elif hasattr(model, 'context_length') and model.context_length:
                    context_size = model.context_length
                elif hasattr(model, 'max_context_length') and model.max_context_length:
                    context_size = model.max_context_length
                
                # Extract pricing if available (OpenRouter-style)
                pricing = None
                if hasattr(model, 'pricing') and model.pricing:
                    pricing = model.pricing
                elif hasattr(model, 'top_provider') and model.top_provider:
                    # Try to extract from top_provider
                    top_provider = model.top_provider
                    if hasattr(top_provider, 'dict'):
                        top_provider = top_provider.dict()
                    if isinstance(top_provider, dict):
                        # Check for pricing in top_provider
                        tp_pricing = top_provider.get('pricing')
                        if tp_pricing:
                            pricing = tp_pricing
                
                result.append(Model(
                    id=model.id,
                    name=model.id,
                    provider_id=self.provider_id,
                    context_size=context_size,
                    context_length=context_size,
                    pricing=pricing
                ))
            
            return result
        except Exception as e:
            import logging
            logging.error(f"OpenAIProviderHandler: Error getting models: {str(e)}", exc_info=True)
            raise e
