import gradio as gr
import os
from groq import Groq
from functools import lru_cache
import time
import requests
import json

# Initialize Groq client
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# Load Army Doctrine snippets
with open('army_doctrine_snippets.json', 'r') as f:
    DOCTRINE_SNIPPETS = json.load(f)['snippets']

@lru_cache(maxsize=100)
def cached_generate_response(query):
    # Prepare a context with relevant doctrine snippets
    relevant_snippets = [snippet for snippet in DOCTRINE_SNIPPETS if any(keyword in query.lower() for keyword in snippet['keywords'])]
    context = "\n".join([snippet['text'] for snippet in relevant_snippets[:3]])  # Use top 3 relevant snippets

    prompt = f"""You are an AI assistant specialized in U.S. Army Doctrine, designed to assist soldiers in the field. 
    Provide a concise, accurate, and actionable response to the following query, based on official U.S. Army Doctrine.
    Your response should be clear, direct, and suitable for field conditions.
    
    Relevant Doctrine Context:
    {context}
    
    Soldier's Query: {query}
    
    Response:"""

    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a U.S. Army Doctrine expert assistant. Provide clear, concise answers based on official doctrine.",
                },
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model="mixtral-8x7b-32768",
            temperature=0.2,  # Lower temperature for more consistent, doctrine-based responses
            max_tokens=200,
        )
        return chat_completion.choices[0].message.content
    except requests.exceptions.RequestException as e:
        return f"Network error: {str(e)}"
    except Exception as e:
        return f"An error occurred: {str(e)}"

def generate_response(query):
    max_retries = 3
    retry_delay = 1  # second

    for attempt in range(max_retries):
        try:
            response = cached_generate_response(query)
            return response + "\n\nNote: This is AI-generated guidance based on U.S. Army Doctrine. Always verify with official sources and commanding officers."
        except Exception as e:
            if attempt < max_retries - 1:
                time.sleep(retry_delay)
                retry_delay *= 2  # Exponential backoff
            else:
                return f"Failed to generate response after {max_retries} attempts. Error: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Enter your field situation or doctrine query here..."),
    outputs="text",
    title="U.S. Army Doctrine Field Assistant",
    description="Get quick, doctrine-based guidance for field situations. Always verify with official sources and commanding officers.",
)

# Launch the app
iface.launch()