import gradio as gr
from ctransformers import AutoModelForCausalLM
import datetime
import json
import os
import psutil
import GPUtil

# Create logs directory if it doesn't exist
if not os.path.exists('logs'):
    os.makedirs('logs')

class SystemMonitor:
    def __init__(self):
        self.gpu = None
        try:
            self.gpu = GPUtil.getGPUs()[0]
        except:
            pass

    def get_stats(self):
        stats = {
            'cpu_percent': psutil.cpu_percent(),
            'memory_percent': psutil.virtual_memory().percent,
            'gpu_util': 0,
            'gpu_memory': 0
        }
        if self.gpu:
            stats['gpu_util'] = self.gpu.load * 100
            stats['gpu_memory'] = self.gpu.memoryUtil * 100
        return stats

class SessionLogger:
    def __init__(self):
        self.session_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_file = f"logs/session_{self.session_id}.json"
        self.session_data = {
            "session_id": self.session_id,
            "start_time": datetime.datetime.now().isoformat(),
            "conversations": []
        }

    def log_interaction(self, user_message, ai_response, stats):
        interaction = {
            "timestamp": datetime.datetime.now().isoformat(),
            "user_message": user_message,
            "ai_response": ai_response,
            "system_stats": stats
        }
        self.session_data["conversations"].append(interaction)
        self.save_session()

    def save_session(self):
        with open(self.log_file, 'w', encoding='utf-8') as f:
            json.dump(self.session_data, f, indent=2, ensure_ascii=False)

# Initialize monitors
system_monitor = SystemMonitor()
session_logger = SessionLogger()

# Initialize model
MODEL_PATH = "D:/llama-tutor/models/llama-2-7b-chat.gguf"
llm = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    model_type='llama',
    context_length=2048,
    gpu_layers=32,
    threads=6,
    batch_size=256
)

SYSTEM_PROMPT = """You are an expert AI tutor, designed to help students learn and master any subject through engaging, Socratic dialogue. Your approach is guided by these core principles:

1. TEACHING APPROACH:
- Focus on building deep understanding rather than providing direct answers
- Use the Socratic method: guide students through questions and critical thinking
- Adapt your explanations to match the student's current level of understanding
- Be patient, encouraging, and maintain a supportive learning environment"""

def manage_conversation_length(history, max_length=4000):
    """Manage conversation history to prevent context overflow"""
    current_length = 0
    truncated_history = []

    # Start from the most recent messages
    for h in reversed(history):
        # Estimate tokens (roughly 4 chars per token)
        message_length = len(h[0]) + len(h[1])
        if current_length + message_length < max_length:
            truncated_history.insert(0, h)
            current_length += message_length
        else:
            # If including the next message would exceed the max length,
            # truncate the current message to fit within the remaining space
            remaining_space = max_length - current_length
            if remaining_space > 0:
                truncated_message = h[1][:remaining_space]
                truncated_history.insert(0, [h[0], truncated_message])
            break

    return truncated_history

def generate_response(message, history):
    # Manage conversation length
    managed_history = manage_conversation_length(history)
    
    # Format the conversation history
    chat_history = "\n".join([f"Human: {h}\nAssistant: {a}\n" for h, a in managed_history])
    
    # Calculate approximate token length
    total_length = len(SYSTEM_PROMPT) + len(chat_history) + len(message)
    approximate_tokens = total_length // 4  # Rough estimation
    
    prompt = f"""<s>[INST] <<SYS>>{SYSTEM_PROMPT}<</SYS>>

Previous conversation:
{chat_history}

Current question: {message} [/INST]"""
    
    # Get system stats before processing
    stats_before = system_monitor.get_stats()
    
    # Update stats with token information
    stats_before['approximate_tokens'] = approximate_tokens
    
    response = llm(prompt, 
              max_new_tokens=1024,  # Increased from 512
              temperature=0.7,
              top_p=0.9,
              top_k=40,
              stop=["</s>"])
    
    # Log interaction with system stats
    session_logger.log_interaction(message, response, stats_before)
    
    return response

def update_stats():
    stats = system_monitor.get_stats()
    current_tokens = len(SYSTEM_PROMPT) // 4  # Base system prompt tokens
    return f"""System Statistics:
    CPU Usage: {stats['cpu_percent']:.1f}%
    Memory Usage: {stats['memory_percent']:.1f}%
    GPU Utilization: {stats['gpu_util']:.1f}%
    GPU Memory: {stats['gpu_memory']:.1f}%
    Context Usage: ~{current_tokens} tokens"""

# Create the Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎓 TA Tutor")
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                value=[],
                label="Tutoring Session",
                height=380
            )
            with gr.Row():
                msg = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask your question here...",
                    lines=2
                )
                submit = gr.Button("Send", variant="primary")
            
            example_questions = [
                "Can you help me understand quantum mechanics?",
                "What's the best way to learn calculus?",
                "Explain photosynthesis in simple terms.",
                "Can you create a study plan for me?",
                "A car accelerates uniformly from rest to a speed of 20 m/s in 5 seconds. What is the acceleration of the car?",
                "Balance the following chemical equation: Fe + O2 -> Fe2O3",
                "Find the derivative of the function f(x) = 3x^4 - 2x^3 + 5x - 7.",
                "Given the matrix A = [[1, 2], [3, 4]], find the determinant of A.",
                "Write a Python function that takes a list of integers as input and returns the sum of all even numbers in the list.",
                "Two forces, F1 = 20 N and F2 = 30 N, act on a point at an angle of 60° to each other. Find the magnitude and direction of the resultant force.",
                "In a series circuit with a 12V battery, there are two resistors: R1 = 4Ω and R2 = 6Ω. Find the current flowing through the circuit.",
                "Describe the basic steps in the engineering design process.",
                "Explain the difference between elastic and plastic deformation in materials.",
                "State the first law of thermodynamics and provide an example of its application in engineering."
            ]
            gr.Examples(
                examples=example_questions,
                inputs=msg,
                label="Example Questions"
            )
        
        with gr.Column(scale=1):
            stats = gr.Textbox(
                label="System Performance",
                lines=6,
                interactive=False,
                value=update_stats()
            )

    def respond(message, history):
        bot_message = generate_response(message, history)
        history.append([message, bot_message])
        new_stats = update_stats()
        return "", history, new_stats

    # Connect message input and submit button
    msg.submit(respond, [msg, chatbot], [msg, chatbot, stats])
    submit.click(respond, [msg, chatbot], [msg, chatbot, stats])

    # Add a refresh button for stats
    refresh = gr.Button("🔄 Refresh Stats")
    refresh.click(update_stats, outputs=[stats])

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_port=7860,
        server_name="0.0.0.0"
    )