import os
import json
import gradio as gr
from datetime import datetime
from huggingface_hub import InferenceClient

# Constants
DEFAULT_MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.1"
HF_TOKEN = os.environ.get("HF_TOKEN")
LOG_FILE = "chat_log.json"

def chat_with_model(system_prompt, user_message, max_tokens=500, model_id=""):
    """
    Generate a chat completion using the specified or default Mistral model.
    Args:
        system_prompt (str): The system prompt to set the context.
        user_message (str): The user's input message.
        max_tokens (int): Maximum number of tokens to generate.
        model_id (str): The model ID to use. If empty, uses the default model.
    Returns:
        str: The model's response.
    """
    model_id = model_id.strip() if model_id else DEFAULT_MODEL_ID
    client = InferenceClient(model_id, token=HF_TOKEN)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message}
    ]
    response = ""
    try:
        for message in client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True
        ):
            response += message.choices[0].delta.content or ""
    except Exception as e:
        response = f"Error: {str(e)}"
    
    # Log the input and output
    log_chat(system_prompt, user_message, response, max_tokens, model_id)
    
    return response

def log_chat(system_prompt, user_message, response, max_tokens, model_id):
    """
    Log the chat details to a JSON file.
    """
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "system_prompt": system_prompt,
        "user_message": user_message,
        "response": response,
        "max_tokens": max_tokens,
        "model_id": model_id
    }
    
    try:
        # Read existing log file if it exists
        if os.path.exists(LOG_FILE):
            with open(LOG_FILE, 'r') as f:
                log_data = json.load(f)
        else:
            log_data = []
        
        # Append new entry
        log_data.append(log_entry)
        
        # Write updated log back to file
        with open(LOG_FILE, 'w') as f:
            json.dump(log_data, f, indent=2)
    except Exception as e:
        print(f"Error logging chat: {str(e)}")

def create_gradio_interface():
    """Create and configure the Gradio interface."""
    return gr.Interface(
        fn=chat_with_model,
        inputs=[
            gr.Textbox(label="System Prompt", placeholder="Enter the system prompt here..."),
            gr.Textbox(label="User Message", placeholder="Ask a question..."),
            gr.Slider(minimum=50, maximum=1000, value=500, step=50, label="Max Tokens"),
            gr.Textbox(label="Model ID", placeholder=f"Enter model ID (default: {DEFAULT_MODEL_ID})")
        ],
        outputs=gr.Textbox(label="Response"),
        title="Mistral Chatbot",
        description="Chat with Mistral model using your own system prompts and choose your model.",
        examples=[
            ["You are a helpful AI assistant.", "What is the capital of France?", 500, ""],
            ["You are an expert in Python programming.", "Explain list comprehensions.", 500, ""]
        ]
    )

if __name__ == "__main__":
    iface = create_gradio_interface()
    iface.launch(show_api=True, share=False)