import gradio as gr
from ctransformers import AutoModelForCausalLM
import datetime
import json
import os
import psutil
import GPUtil
import torch
from typing import Dict, Tuple, Any, Optional
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/app.log'),
        logging.StreamHandler()
    ]
)

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

class EnhancedSystemMonitor:
    """Enhanced system monitoring with detailed GPU stats and change tracking"""
    def __init__(self):
        self.gpu = None
        self.previous_stats = None
        try:
            self.gpu = GPUtil.getGPUs()[0]
            logging.info(f"GPU initialized: {self.gpu.name}")
        except Exception as e:
            logging.warning(f"Could not initialize GPU monitoring: {e}")

    def get_stats(self) -> Dict[str, float]:
        """Get comprehensive system statistics including GPU metrics"""
        stats = {
            'cpu_percent': psutil.cpu_percent(),
            'memory_percent': psutil.virtual_memory().percent,
            'gpu_util': 0,
            'gpu_memory': 0,
            'gpu_clock': 0,
            'gpu_temp': 0,
            'gpu_power': 0
        }
        
        if self.gpu:
            try:
                self.gpu = GPUtil.getGPUs()[0]  # Refresh GPU info
                stats.update({
                    'gpu_util': self.gpu.load * 100,
                    'gpu_memory': self.gpu.memoryUtil * 100,
                    'gpu_clock': getattr(self.gpu, 'memoryTotal', 0),  # Using memoryTotal as a fallback
                    'gpu_temp': getattr(self.gpu, 'temperature', 0),
                    'gpu_power': getattr(self.gpu, 'powerDraw', 0)
                })
            except Exception as e:
                logging.error(f"Error updating GPU stats: {e}")

        if self.previous_stats:
            stats['gpu_util_change'] = stats['gpu_util'] - self.previous_stats['gpu_util']
            stats['gpu_memory_change'] = stats['gpu_memory'] - self.previous_stats['gpu_memory']
        
        self.previous_stats = stats.copy()
        return stats

class SessionLogger:
    """Enhanced session logging with system metrics"""
    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(),
            "system_info": self._get_system_info(),
            "conversations": []
        }

    def _get_system_info(self) -> Dict[str, Any]:
        """Collect system information at session start"""
        info = {
            "cpu_freq": psutil.cpu_freq()._asdict() if psutil.cpu_freq() else {},
            "cpu_count": psutil.cpu_count(),
            "cpu_count_logical": psutil.cpu_count(logical=True),
            "memory_total": psutil.virtual_memory().total,
            "platform": os.name
        }
        try:
            gpu = GPUtil.getGPUs()[0]
            info["gpu_info"] = {
                "name": gpu.name,
                "memory_total": gpu.memoryTotal,
                "driver": gpu.driver
            }
        except:
            info["gpu_info"] = None
        return info

    def log_interaction(self, 
                       user_message: str, 
                       ai_response: str, 
                       stats: Dict[str, float],
                       metadata: Optional[Dict[str, Any]] = None):
        """Log interaction with enhanced metadata"""
        interaction = {
            "timestamp": datetime.datetime.now().isoformat(),
            "user_message": user_message,
            "ai_response": ai_response,
            "system_stats": stats,
            "metadata": metadata or {}
        }
        self.session_data["conversations"].append(interaction)
        self.save_session()

    def save_session(self):
        """Save session data to file with error handling"""
        try:
            with open(self.log_file, 'w', encoding='utf-8') as f:
                json.dump(self.session_data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            logging.error(f"Error saving session data: {e}")

class ModelManager:
    """Manages model initialization, inference, and resource handling"""
    def __init__(self, model_path: str):
        self.model_path = model_path
        self.llm = None
        self.config = None
        self.initialize_model()
        self.total_tokens_processed = 0

    def initialize_model(self, gpu_layers: int = 24, batch_size: int = 128) -> None:
        """Initialize the model with specified parameters"""
        try:
            self.llm = AutoModelForCausalLM.from_pretrained(
                self.model_path,
                model_type='llama',
                context_length=2048,
                gpu_layers=gpu_layers,
                threads=4,
                batch_size=batch_size
            )
            self.config = {'gpu_layers': gpu_layers, 'batch_size': batch_size}
            logging.info(f"Model initialized with config: {self.config}")
        except Exception as e:
            logging.error(f"Error initializing model: {e}")
            raise

    def generate_with_monitoring(self, 
                               prompt: str, 
                               system_monitor: EnhancedSystemMonitor,
                               **kwargs) -> Tuple[str, Dict[str, Any]]:
        """Generate response with comprehensive monitoring"""
        pre_stats = system_monitor.get_stats()
        
        try:
            # Generate response
            response = self.llm(prompt, **kwargs)
            post_stats = system_monitor.get_stats()
            
            # Calculate tokens for this interaction
            input_tokens = len(prompt) // 4
            output_tokens = len(response) // 4
            total_tokens = input_tokens + output_tokens
            
            # Update total tokens processed
            self.total_tokens_processed += total_tokens
            
            # Update stats with token information
            post_stats.update({
                'approximate_tokens': total_tokens,
                'total_tokens_processed': self.total_tokens_processed,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens
            })
            
            return response, {
                'pre_stats': pre_stats,
                'post_stats': post_stats,
                'tokens_processed': total_tokens,
                'total_tokens': self.total_tokens_processed
            }
            
        except RuntimeError as e:
            if "out of memory" in str(e):
                logging.warning("OOM detected, attempting recovery...")
                torch.cuda.empty_cache()
                
                # Reduce parameters and retry
                new_config = {
                    'gpu_layers': max(8, self.config['gpu_layers'] - 4),
                    'batch_size': max(32, self.config['batch_size'] // 2)
                }
                
                self.initialize_model(**new_config)
                return self.generate_with_monitoring(prompt, system_monitor, **kwargs)
            raise

def format_stats(stats: Dict[str, float]) -> str:
    """Format system statistics for display"""
    return f"""System Statistics:
    CPU Usage: {stats.get('cpu_percent', 0):.1f}%
    Memory Usage: {stats.get('memory_percent', 0):.1f}%
    
    GPU Status:
    • Utilization: {stats.get('gpu_util', 0):.1f}% ({stats.get('gpu_util_change', 0):.1f}% change)
    • Memory: {stats.get('gpu_memory', 0):.1f}% ({stats.get('gpu_memory_change', 0):.1f}% change)
    • Memory Total: {stats.get('gpu_clock', 0)} MB
    • Temperature: {stats.get('gpu_temp', 0)}°C
    
    Model Info:
    • Current Usage: ~{stats.get('approximate_tokens', 0)} tokens
    • Input Tokens: ~{stats.get('input_tokens', 0)}
    • Output Tokens: ~{stats.get('output_tokens', 0)}
    • Total Processed: {stats.get('total_tokens_processed', 0)} tokens"""

def create_demo(model_path: str):
    """Create and configure the Gradio demo"""
    # Initialize components
    system_monitor = EnhancedSystemMonitor()
    session_logger = SessionLogger()
    model_manager = ModelManager(model_path)

    # System prompt
    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: int = 4000) -> list:
        """Manage conversation history to prevent context overflow"""
        current_length = 0
        truncated_history = []
        
        for h in reversed(history):
            message_length = len(h['content'])
            if current_length + message_length < max_length:
                truncated_history.insert(0, h)
                current_length += message_length
            else:
                remaining_space = max_length - current_length
                if remaining_space > 0:
                    truncated_message = h.copy()
                    truncated_message['content'] = h['content'][:remaining_space]
                    truncated_history.insert(0, truncated_message)
                break
        
        return truncated_history

    def generate_response(message: str, history: list) -> Tuple[str, Dict[str, Any]]:
        """Generate response with conversation management"""
        managed_history = manage_conversation_length(history)
        
        chat_history = "\n".join([
            f"{'Human' if h['role'] == 'user' else 'Assistant'}: {h['content']}"
            for h in managed_history
        ])
        
        prompt = f"""<s>[INST] <<SYS>>{SYSTEM_PROMPT}<</SYS>>

        Previous conversation:
        {chat_history}

        Current question: {message} [/INST]"""
        
        # Calculate approximate tokens for the full context
        system_tokens = len(SYSTEM_PROMPT) // 4
        history_tokens = len(chat_history) // 4
        message_tokens = len(message) // 4
        total_input_tokens = system_tokens + history_tokens + message_tokens
        
        # Update system monitor with token count
        current_stats = system_monitor.get_stats()
        current_stats['approximate_tokens'] = total_input_tokens
        
        response, monitoring_data = model_manager.generate_with_monitoring(
            prompt,
            system_monitor,
            max_new_tokens=1024,
            temperature=0.7,
            top_p=0.9,
            top_k=40,
            stop=["</s>"]
        )
        
        # Update monitoring data with token information
        monitoring_data['post_stats']['approximate_tokens'] = total_input_tokens + (len(response) // 4)
        
        session_logger.log_interaction(
            message,
            response,
            monitoring_data['post_stats'],
            {"monitoring_data": monitoring_data}
        )
        
        return response, monitoring_data

    # 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,
                    type="messages"  # Updated format
                )
                
                with gr.Row():
                    msg = gr.Textbox(
                        label="Your Question",
                        placeholder="Ask your question here...",
                        lines=2
                    )
                    submit = gr.Button("Send", variant="primary")
                
                # Example questions
                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_display = gr.Textbox(
                    label="System Performance",
                    lines=8,
                    interactive=False,
                    value=format_stats(system_monitor.get_stats())
                )
        
        def respond(message: str, history: list) -> Tuple[str, list, str]:
            """Handle user input and generate response"""
            try:
                response, monitoring_data = generate_response(message, history)
                
                history.append({"role": "user", "content": message})
                history.append({"role": "assistant", "content": response})
                
                new_stats = format_stats(monitoring_data['post_stats'])
                return "", history, new_stats
                
            except Exception as e:
                error_message = f"Error during generation: {str(e)}"
                logging.error(error_message)
                history.append({"role": "system", "content": error_message})
                return "", history, format_stats(system_monitor.get_stats())

        # Connect UI elements
        msg.submit(respond, [msg, chatbot], [msg, chatbot, stats_display])
        submit.click(respond, [msg, chatbot], [msg, chatbot, stats_display])

        # Add stats refresh button
        refresh = gr.Button("🔄 Refresh Stats")
        refresh.click(
            lambda: format_stats(system_monitor.get_stats()),
            outputs=[stats_display]
        )

    return demo

if __name__ == "__main__":
    MODEL_PATH = "D:/llama-tutor/models/llama-2-7b-chat.gguf"
    demo = create_demo(MODEL_PATH)
    demo.queue()
    demo.launch(
        server_port=7860,
        server_name="127.0.0.1",  # Changed to localhost for security
        show_error=True
    )