import os
import requests

class ChatAI:
    """Simple wrapper supporting OpenAI and Google Gemini HTTP calls.
    This class performs minimal error handling; adapt/log as needed.
    """
    def __init__(self, provider="openai", api_key=None, model="gpt-4o-mini"):
        self.provider = provider.lower()
        self.api_key = api_key or os.getenv('AYUAI_OPENAI_KEY') or os.getenv('AYUAI_KEY')
        self.model = model

    def ask(self, prompt):
        if self.provider == 'openai':
            return self._ask_openai(prompt)
        elif self.provider == 'gemini':
            return self._ask_gemini(prompt)
        else:
            raise ValueError('Unsupported provider: ' + str(self.provider))

    def _ask_openai(self, prompt):
        """Call OpenAI Chat Completions (HTTP). Requires valid API key in self.api_key"""
        if not self.api_key:
            raise ValueError('OpenAI API key missing. Set api_key or AYUAI_OPENAI_KEY env var.')

        url = "https://api.openai.com/v1/chat/completions"
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        data = {
            'model': self.model,
            'messages': [{'role': 'user', 'content': prompt}]
        }
        resp = requests.post(url, headers=headers, json=data, timeout=30)
        resp.raise_for_status()
        jr = resp.json()
        # Defensive access
        return jr.get('choices', [{}])[0].get('message', {}).get('content', '')

    def _ask_gemini(self, prompt):
        """Call Google Generative API (Gemini) - uses API key in URL query param by default."""
        if not self.api_key:
            raise ValueError('Google API key missing. Set api_key or AYUAI_GOOGLE_KEY env var.')

        url = f"https://generativelanguage.googleapis.com/v1beta2/models/{self.model}:generateText?key={self.api_key}"
        # Note: Google has different endpoints / versions — update model string & endpoint as needed.
        data = {
            'prompt': {
                'text': prompt
            },
            'maxOutputTokens': 512
        }
        resp = requests.post(url, json=data, timeout=30)
        resp.raise_for_status()
        jr = resp.json()
        # This may vary depending on API response format; attempt multiple fallbacks
        if 'candidates' in jr and jr['candidates']:
            return jr['candidates'][0].get('output', '')
        return jr.get('output', '') or str(jr)
