import gradio as gr
from huggingface_hub import InferenceClient
from textblob import TextBlob
import json
import os
import time
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)

# Get the API token from the environment variable
api_token = os.getenv('HUGGINGFACEHUB_API_TOKEN')

client = InferenceClient(
    model="Futuresony/future_ai_12_10_2024.gguf",
    token=api_token
)

# Directory to store interactions and feedback
DATA_DIR = "data"
INTERACTIONS_FILE = os.path.join(DATA_DIR, "interactions.json")

# Ensure the data directory exists
os.makedirs(DATA_DIR, exist_ok=True)

def format_alpaca_prompt(user_input, system_prompt, history):
    """Formats input in Alpaca/LLaMA style"""
    history_str = "\n".join([f"### Instruction:\n{h[0]}\n### Response:\n{h[1]}" for h in history])
    prompt = f"""{system_prompt}
{history_str}

### Instruction:
{user_input}

### Response:
"""
    return prompt

def analyze_sentiment(message):
    """Analyze the sentiment of the user's message"""
    blob = TextBlob(message)
    sentiment = blob.sentiment.polarity
    return sentiment

def save_interaction(user_input, chatbot_response, feedback=None):
    """Save the interaction and feedback to a file"""
    interaction = {
        "user_input": user_input,
        "chatbot_response": chatbot_response,
        "feedback": feedback,
        "timestamp": "2025-02-25 14:08:31"
    }
    if os.path.exists(INTERACTIONS_FILE):
        with open(INTERACTIONS_FILE, "r") as file:
            interactions = json.load(file)
    else:
        interactions = []

    interactions.append(interaction)

    with open(INTERACTIONS_FILE, "w") as file:
        json.dump(interactions, file, indent=4)

def respond(message, history, system_message, max_tokens, temperature, top_p, feedback=None):
    sentiment = analyze_sentiment(message)

    # Adjust system message based on sentiment
    if sentiment < -0.2:
        system_message = "You are a sympathetic Chatbot."
    elif sentiment > 0.2:
        system_message = "You are an enthusiastic Chatbot."
    else:
        system_message = "You are a friendly Chatbot."

    formatted_prompt = format_alpaca_prompt(message, system_message, history)

    # Retry mechanism
    max_retries = 3
    for attempt in range(max_retries):
        try:
            response = client.text_generation(
                formatted_prompt,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
            )
            break  # Exit the loop if the request is successful
        except Exception as e:
            logging.error(f"Attempt {attempt + 1} failed: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # Exponential backoff
            else:
                raise e

    # ✅ Extract only the response
    cleaned_response = response.split("### Response:")[-1].strip()

    history.append((message, cleaned_response))  # ✅ Update history with the new message and response

    save_interaction(message, cleaned_response, feedback)  # ✅ Save the interaction and feedback

    yield cleaned_response  # ✅ Output only the answer

def collect_feedback(response, feedback):
    """Collect user feedback on the chatbot's response"""
    save_interaction(response, feedback=feedback)

def view_interactions():
    if os.path.exists(INTERACTIONS_FILE):
        with open(INTERACTIONS_FILE, "r") as file:
            interactions = json.load(file)
            return json.dumps(interactions, indent=4)
    else:
        return "No interactions found."

def download_interactions():
    """Provide the interactions file for download"""
    if os.path.exists(INTERACTIONS_FILE):
        return INTERACTIONS_FILE
    else:
        return None

# Create a Gradio interface to display interactions
view_interface = gr.Interface(
    fn=view_interactions,
    inputs=[],
    outputs="text",
    title="View Interactions"
)

# Create a Gradio interface for downloading interactions
download_interface = gr.Interface(
    fn=download_interactions,
    inputs=[],
    outputs=gr.File(label="Download Interactions"),
    title="Download Interactions"
)

feedback_interface = gr.Interface(
    fn=collect_feedback,
    inputs=[
        gr.Textbox(label="Response"),
        gr.Radio(choices=["Good", "Bad"], label="Feedback"),
    ],
    outputs="text",
    title="Feedback Interface"
)

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=250, value=128, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.9, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.99, step=0.01, label="Top-p (nucleus sampling)"),
    ],
)

if __name__ == "__main__":
    demo.launch()
    feedback_interface.launch()
    view_interface.launch()
    download_interface.launch()