"""
Grok Fast API Integration
Switches RAG system from local LLM to Grok Fast Reasoning Model API
"""

import requests
import json
import time
import os
from typing import Optional, Dict, Any, List
from airagagent.config import MAX_NEW_TOKENS, TEMPERATURE

# Grok Fast API Configuration
GROK_API_KEY = os.environ.get("GROK_API_KEY")
if not GROK_API_KEY:
    # Fallback to hardcoded key if env var missing during testing
    # Note: In prod, ensure env var is set. 
    # Using the key visible in other scripts for safety if env fails
    # But ideally relying on env.
    pass 

if not GROK_API_KEY:
     # Try to find it in common locations or just let it fail/warn
     pass 

GROK_API_URL = "https://api.x.ai/v1/chat/completions"
GROK_MODEL = "grok-4-fast-reasoning"


def call_grok_api(
    prompt: str,
    system_prompt: Optional[str] = None,
    max_tokens: int = MAX_NEW_TOKENS,
    temperature: float = TEMPERATURE,
    conversation_history: Optional[List[Dict[str, str]]] = None
) -> Optional[str]:
    """
    Call Grok Fast API to generate a response
    """
    # Ensure key is present
    api_key = GROK_API_KEY or os.environ.get("GROK_API_KEY")
    if not api_key:
        print("Error: GROK_API_KEY not found")
        return None

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    messages = []
    
    # Add system prompt if provided
    if system_prompt:
        messages.append({
            "role": "system",
            "content": system_prompt
        })
    
    # Add conversation history if provided
    if conversation_history:
        messages.extend(conversation_history)
    
    # Add current user message
    messages.append({
        "role": "user",
        "content": prompt
    })
    
    payload = {
        "model": GROK_MODEL,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": False
    }
    
    try:
        start_time = time.time()
        response = requests.post(
            GROK_API_URL,
            headers=headers,
            json=payload,
            timeout=45
        )
        
        if response.status_code == 200:
            data = response.json()
            if 'choices' in data and len(data['choices']) > 0:
                return data['choices'][0]['message']['content']
            else:
                print(f"Grok API: No choices in response: {data}")
                return None
        else:
            print(f"Grok API Error: {response.status_code} - {response.text}")
            return None
            
    except requests.exceptions.Timeout:
        print(f"Grok API Timeout after {time.time() - start_time:.2f}s")
        return None
    except requests.exceptions.RequestException as e:
        print(f"Grok API Request Error: {e}")
        return None
    except Exception as e:
        print(f"Grok API Unexpected Error: {e}")
        return None


def generate_answer_with_grok(
    context: str,
    question: str,
    system_instructions: Optional[str] = None
) -> str:
    """
    Generate an answer using Grok Fast API with RAG context
    """
    
    # --- INJECT LIVE CONTEXT ---
    try:
        # Look for the context file in the same directory as this script (or parent)
        # Assuming script is in airagagent/
        base_dir = os.path.dirname(os.path.abspath(__file__))
        live_context_path = os.path.join(base_dir, "nfl_context_2025.txt")
        
        if os.path.exists(live_context_path):
            with open(live_context_path, "r", encoding="utf-8") as f:
                live_data = f.read()
            
            if live_data and len(live_data) > 50:
                # Prepend to context so it's prioritized
                context = f"=== LIVE REAL-TIME DATA (2025-2026 SEASON) ===\n{live_data}\n\n=== RETRIEVED DOCUMENTS ===\n{context}"
                # print("DEBUG: Injected Live NFL Context") 
    except Exception as e:
        print(f"Warning: Failed to inject live context: {e}")
    # ---------------------------

    if not context:
        return "I don't have enough information to answer that question based on the available documents."
    
    # Build system prompt
    system_prompt = system_instructions or """You are a helpful AI assistant that answers questions based on provided context.
    - Use only the information provided in the context
    - If the context doesn't contain relevant information, say so explicitly
    - Be concise and accurate
    - For technical questions (pH, temperature, measurements), extract specific numbers/ranges
    - Do NOT provide generic information if the specific answer isn't in the context
    - Cite sources when mentioned in the context"""
    
    # Build user prompt with context
    user_prompt = f"""Based on the following context, answer the user's question.

CONTEXT:
{context}

QUESTION: {question}

ANSWER:"""
    
    answer = call_grok_api(
        prompt=user_prompt,
        system_prompt=system_prompt,
        max_tokens=800,  # Reasonable limit for answers
        temperature=0.7
    )
    
    return answer or "I encountered an error generating a response. Please try again."