import gradio as gr
from transformers import pipeline
import uuid

# Use Hugging Face hosted pipeline (won't need to load heavy model)
generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.1")

# Session state store
user_sessions = {}

def get_session(session_id):
    if session_id not in user_sessions:
        user_sessions[session_id] = {"chat": []}
    return user_sessions[session_id]

def build_prompt(history):
    prompt = "You are a helpful assistant.\n"
    for turn in history:
        prompt += f"User: {turn[0]}\nAssistant: {turn[1]}\n"
    return prompt

def chat_fn(user_input, session_id):
    session = get_session(session_id)
    history = session["chat"]
    prompt = build_prompt(history + [[user_input, ""]]) + "Assistant:"
    result = generator(prompt, max_new_tokens=128, do_sample=True, temperature=0.7)[0]["generated_text"]
    reply = result[len(prompt):].strip().split("\n")[0]
    history.append([user_input, reply])
    return history, session_id

def upload_handler(file):
    try:
        content = file.read().decode("utf-8")
        return content
    except Exception as e:
        return f"Error: {str(e)}"

def code_viewer(code):
    # Simulate code output
    return "Code received. Execution is disabled on this free demo for safety."

# Create unique session ID
def create_session():
    return str(uuid.uuid4())

with gr.Blocks(title="Multi-Tool AI Chatbot") as demo:
    session_id_state = gr.State(value=create_session())

    with gr.Tab("Chatbot"):
        chatbot = gr.Chatbot()
        msg = gr.Textbox(label="Your message")
        send = gr.Button("Send")
        send.click(
            fn=lambda msg, sid: chat_fn(msg, sid),
            inputs=[msg, session_id_state],
            outputs=[chatbot, session_id_state]
        )

    with gr.Tab("File Reader"):
        file_input = gr.File(label="Upload .txt file", file_types=[".txt"])
        file_output = gr.Textbox(label="File content")
        file_input.change(upload_handler, inputs=file_input, outputs=file_output)

    with gr.Tab("Code Viewer"):
        code_input = gr.Code(language="python", label="Enter Python code")
        code_output = gr.Textbox(label="Output")
        run_btn = gr.Button("View Output")
        run_btn.click(fn=code_viewer, inputs=code_input, outputs=code_output)

if __name__ == "__main__":
    demo.launch()