import subprocess
import sys
import os
import time
import json
from datetime import datetime
from IPython.display import display, HTML, clear_output, Markdown

# Install required packages
print("📦 Installing required packages...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "openai"])
print("✅ Packages installed successfully!")

from openai import OpenAI

def display_status(message, status="info"):
    """Display colored status messages"""
    colors = {
        "info": "#3498db",
        "success": "#2ecc71",
        "warning": "#f39c12",
        "error": "#e74c3c",
        "processing": "#9b59b6"
    }
    html = f"""
    <div style="padding: 10px; margin: 10px 0; border-left: 4px solid {colors.get(status, '#3498db')}; background-color: #f8f9fa;">
        <strong style="color: {colors.get(status, '#3498db')};">{message}</strong>
    </div>
    """
    display(HTML(html))

def display_progress(current, total, label="Progress"):
    """Display a progress bar"""
    percentage = (current / total) * 100
    bar_length = 50
    filled_length = int(bar_length * current // total)
    bar = '█' * filled_length + '░' * (bar_length - filled_length)
  
    html = f"""
    <div style="margin: 10px 0;">
        <div style="font-weight: bold; margin-bottom: 5px;">{label}</div>
        <div style="background-color: #ecf0f1; border-radius: 10px; padding: 3px;">
            <div style="background-color: #3498db; width: {percentage}%; border-radius: 10px; padding: 5px; color: white; text-align: center;">
                {bar} {percentage:.1f}%
            </div>
        </div>
    </div>
    """
    display(HTML(html))

display_status("🚀 Setting up Ollama...", "processing")

# Install Ollama
print("Installing Ollama... This may take a minute...")
result = os.system("curl -fsSL https://ollama.com/install.sh | sh 2>/dev/null")
if result == 0:
    display_status("✅ Ollama installed successfully!", "success")
else:
    display_status("⚠️ Ollama installation had warnings but may still work", "warning")

# Start Ollama server
print("Starting Ollama server...")
os.system("nohup ollama serve > /tmp/ollama_serve_stdout.log 2>/tmp/ollama_serve_stderr.log &")
time.sleep(5)

# Check if running
running = os.system("ps aux | grep -E 'ollama serve' | grep -v grep > /dev/null 2>&1")
if running == 0:
    display_status("✅ Ollama server is running!", "success")
else:
    display_status("❌ Ollama server failed to start. Check troubleshooting section.", "error")

display_status("📥 Downloading GPT-OSS:20B Model (~13GB)", "processing")
print("This will take several minutes. Please be patient...")
print("="*60)

# Start the download
start_time = time.time()
result = os.system("ollama pull gpt-oss:20b")
end_time = time.time()

if result == 0:
    elapsed = end_time - start_time
    display_status(f"✅ Model downloaded successfully in {elapsed/60:.1f} minutes!", "success")
else:
    display_status("❌ Model download failed. Please check your connection and try again.", "error")

# Verify model is available
print("\n📋 Available models:")
os.system("ollama list")

class GPTOSSChat:
    def __init__(self, system_message = "You are a helpful AI assistant."):
        self.client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")
        self.conversation_history = []
        self.system_message = system_message
      
    def set_system_message(self, message):
        """Change the system message"""
        self.system_message = message
        display_status(f"System message updated: {message[:100]}...", "info")
  
    def clear_history(self):
        """Clear conversation history"""
        self.conversation_history = []
        display_status("Conversation history cleared", "info")
  
    def chat(self, user_input, include_history=True):
        """Send a message to the model"""
        try:
            messages = [{"role": "system", "content": self.system_message}]
          
            if include_history:
                messages.extend(self.conversation_history)
          
            messages.append({"role": "user", "content": user_input})
          
            # Show processing indicator
            display_status("🤔 Thinking...", "processing")
          
            response = self.client.chat.completions.create(
                model="gpt-oss:20b",
                messages=messages
            )
          
            assistant_response = response.choices[0].message.content
          
            # Update history
            if include_history:
                self.conversation_history.append({"role": "user", "content": user_input})
                self.conversation_history.append({"role": "assistant", "content": assistant_response})
          
            # Display response with nice formatting
            clear_output(wait=True)
            display(Markdown(f"**You:** {user_input}"))
            display(Markdown(f"**Assistant:** {assistant_response}"))
          
            return assistant_response
          
        except Exception as e:
            display_status(f"Error: {str(e)}", "error")
            return None
  
    def interactive_chat(self):
        """Start an interactive chat session"""
        display(HTML("""
        <div style="background-color: #3498db; color: white; padding: 15px; border-radius: 10px; margin: 10px 0;">
            <h3>🤖 Interactive Chat with GPT-OSS:20B</h3>
            <p>Type your messages below. Type 'exit' to end the chat, 'clear' to clear history, or 'system:' followed by a message to change the system prompt.</p>
        </div>
        """))
      
        while True:
            user_input = input("\n💬 You: ")
          
            if user_input.lower() == 'exit':
                display_status("Chat session ended. Goodbye! 👋", "info")
                break
            elif user_input.lower() == 'clear':
                self.clear_history()
                continue
            elif user_input.lower().startswith('system:'):
                new_system = user_input[7:].strip()
                self.set_system_message(new_system)
                continue
          
            self.chat(user_input)

# Initialize the chat interface
display_status("🎉 Initializing GPT-OSS Chat Interface...", "processing")
chat = GPTOSSChat()
pirate_chat = GPTOSSChat(system_message = 'You only reply as a pirate!')
display_status("✅ Chat interface ready!", "success")

display(HTML("<h2>🧪 Testing the Model</h2>"))

# Test with a simple query
test_response = chat.chat("Hello! Can you confirm you're working properly? Please respond with a brief greeting.", include_history=False)

if test_response:
    display_status("✅ Model is working perfectly!", "success")
else:
    display_status("❌ Model test failed. Please check the troubleshooting section.", "error")

# Example 1: Simple query
display(Markdown("### Example 1: Simple Query"))
print("chat.chat('What is the capital of France?')")

# Example 2: Creative writing
display(Markdown("### Example 2: Creative Writing"))
chat.set_system_message("You are a creative writer who specializes in science fiction.")
creative_response = chat.chat("Write a haiku about artificial intelligence", include_history=False)

# Example 3: Code assistance
display(Markdown("### Example 3: Code Assistance"))
chat.set_system_message("You are a helpful coding assistant.")
code_response = chat.chat("Write a Python function to calculate fibonacci numbers", include_history=False)

# Reset to default
chat.set_system_message("You are a helpful AI assistant.")
chat.clear_history()

from __future__ import annotations
from typing import Iterable, List, Sequence

# ------------------------------------------------------------------
# 1️⃣  Simple pure recursion (good for tiny values of n)
# ------------------------------------------------------------------
def fib_recursive(n: int) -> int:
    """Return the *n*‑th Fibonacci number (0‑based).

    >>> fib_recursive(0)
    0
    >>> fib_recursive(1)
    1
    >>> fib_recursive(10)
    55
    """
    if n < 0:
        raise ValueError("n must be a non‑negative integer.")
    if n <= 1:
        return n
    return fib_recursive(n - 1) + fib_recursive(n - 2)


# ------------------------------------------------------------------
# 2️⃣  Recursive + memoisation (fast for repeated calls)
# ------------------------------------------------------------------
def fib_memoised(n: int, cache: dict[int, int] | None = None) -> int:
    """Return the *n*‑th Fibonacci number with memoisation.

    The optional *cache* dict is used internally – you can supply one to
    share memoisation across multiple calls.

    >>> fib_memoised(50)                # 12586269025
    12586269025
    >>> cache: dict[int, int] = {}
    >>> fib_memoised(50, cache=cache)  # uses cache
    12586269025
    """
    if cache is None:
        cache = {}
    if n < 0:
        raise ValueError("n must be a non‑negative integer.")
    if n in cache:
        return cache[n]
    if n <= 1:
        cache[n] = n
    else:
        cache[n] = fib_memoised(n - 1, cache) + fib_memoised(n - 2, cache)
    return cache[n]


# ------------------------------------------------------------------
# 3️⃣  Iterative, constant‑space (fastest for large n)
# ------------------------------------------------------------------
def fib_iterative(n: int) -> int:
    """Return the *n*‑th Fibonacci number using an O(1) space loop.

    >>> fib_iterative(10)
    55
    >>> fib_iterative(100)
    354224848179261915075
    """
    if n < 0:
        raise ValueError("n must be a non‑negative integer.")
    a, b = 0, 1
    for _ in range(n):
        a, b = b, a + b
    return a


# ------------------------------------------------------------------
# 4️⃣  Generator – yields as many Fibonacci numbers as you need
# ------------------------------------------------------------------
def fibonacci_generator() -> Iterable[int]:
    """Yield Fibonacci numbers forever (0, 1, 1, 2, …)."""
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b


# ------------------------------------------------------------------
# 5️⃣  Build a list of the first *n* Fibonacci numbers
# ------------------------------------------------------------------
def first_n_fibonacci(n: int) -> Sequence[int]:
    """Return a list of the first *n* Fibonacci numbers (0‑based).

    >>> first_n_fibonacci(7)
    [0, 1, 1, 2, 3, 5, 8]
    """
    if n < 0:
        raise ValueError("n must be a non‑negative integer.")
    seq: List[int] = []
    for i in range(n):
        seq.append(fib_iterative(i))
    return seq


# ------------------------------------------------------------------
# Quick sanity tests (run `python -m pytest` if you set this up as a test file)
# ------------------------------------------------------------------
if __name__ == "__main__":
    for func in [fib_recursive, fib_memoised, fib_iterative]:
        for n in [0, 1, 5, 10]:
            assert func(n) == fib_iterative(n), f"{func.__name__} failed on n={n}"
    print("All basic checks passed!")

display(HTML("""
<div style="background-color: #2ecc71; color: white; padding: 20px; border-radius: 10px; margin: 20px 0;">
    <h2>💬 Start Interactive Chat</h2>
    <p>Run the cell below to start an interactive chat session with GPT-OSS:20B</p>
    <ul>
        <li>Type your messages and press Enter</li>
        <li>Type 'exit' to end the chat</li>
        <li>Type 'clear' to clear conversation history</li>
        <li>Type 'system: [message]' to change the AI's behavior</li>
    </ul>
</div>
"""))

# Uncomment the line below to start interactive chat
# chat.interactive_chat()

def quick_chat(message, system="You are a helpful AI assistant."):
    """Quick one-off chat without history"""
    try:
        client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")
        response = client.chat.completions.create(
            model="gpt-oss:20b",
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": message}
            ]
        )
        result = response.choices[0].message.content
        display(Markdown(f"**Question:** {message}"))
        display(Markdown(f"**Answer:** {result}"))
        return result
    except Exception as e:
        display_status(f"Error: {str(e)}", "error")
        return None

# Example usage
display(HTML("<h3>Quick Chat Example:</h3>"))
quick_chat("What are the three laws of robotics?")

def check_ollama_status():
    """Check the current status of Ollama"""
    display_status("Checking Ollama status...", "info")
  
    # Check if process is running
    running = os.system("ps aux | grep -E 'ollama serve' | grep -v grep > /dev/null 2>&1")
    if running == 0:
        display_status("✅ Ollama server is running", "success")
    else:
        display_status("❌ Ollama server is not running", "error")
  
    # Check if API is responding
    api_check = os.system("curl -s http://localhost:11434/v1/models > /dev/null 2>&1")
    if api_check == 0:
        display_status("✅ Ollama API is responding", "success")
        print("\nAvailable models:")
        os.system("ollama list")
    else:
        display_status("❌ Ollama API is not responding", "error")

def restart_ollama():
    """Restart Ollama server"""
    display_status("Restarting Ollama server...", "processing")
  
    # Kill existing processes
    os.system("pkill -9 ollama 2>/dev/null")
    time.sleep(2)
  
    # Start server
    os.system("nohup ollama serve > /tmp/ollama_serve_stdout.log 2>/tmp/ollama_serve_stderr.log &")
    time.sleep(5)
  
    # Check status
    check_ollama_status()

def view_ollama_logs():
    """View the last 20 lines of Ollama logs"""
    display_status("📜 Ollama Server Logs (last 20 lines):", "info")
    os.system("tail -20 /tmp/ollama_serve_stdout.log 2>/dev/null")
    os.system("tail -20 /tmp/ollama_serve_stderr.log 2>/dev/null")

# Test the status
check_ollama_status()

def create_chat_widget():
    """Create a simple chat widget for notebook interaction"""
    html_code = """
    <div id="chat-container" style="border: 2px solid #3498db; border-radius: 10px; padding: 20px; margin: 10px 0;">
        <h3>💬 Chat with GPT-OSS:20B</h3>
        <div id="chat-messages" style="height: 300px; overflow-y: auto; border: 1px solid #ecf0f1; padding: 10px; margin: 10px 0; background-color: #f8f9fa;">
            <p style="color: #7f8c8d;">Chat messages will appear here...</p>
        </div>
        <div style="display: flex; gap: 10px;">
            <input type="text" id="user-input" placeholder="Type your message here..." 
                   style="flex: 1; padding: 10px; border: 1px solid #3498db; border-radius: 5px;"
                   onkeypress="if(event.key === 'Enter') sendMessage()">
            <button onclick="sendMessage()" 
                    style="padding: 10px 20px; background-color: #3498db; color: white; border: none; border-radius: 5px; cursor: pointer;">
                Send
            </button>
            <button onclick="clearChat()" 
                    style="padding: 10px 20px; background-color: #e74c3c; color: white; border: none; border-radius: 5px; cursor: pointer;">
                Clear
            </button>
        </div>
    </div>
  
    <script>
    function sendMessage() {
        const input = document.getElementById('user-input');
        const messages = document.getElementById('chat-messages');
      
        if (input.value.trim() === '') return;
      
        // Add user message to chat
        messages.innerHTML += '<div style="margin: 10px 0;"><strong>You:</strong> ' + input.value + '</div>';
      
        // Add processing indicator
        messages.innerHTML += '<div style="margin: 10px 0; color: #7f8c8d;"><em>Processing...</em></div>';
      
        // Note: In a real implementation, you would send this to the kernel
        messages.innerHTML += '<div style="margin: 10px 0; color: #3498db;"><strong>Note:</strong> Run chat.chat("' + input.value + '") in the next cell to get the response.</div>';
      
        input.value = '';
        messages.scrollTop = messages.scrollHeight;
    }
  
    function clearChat() {
        document.getElementById('chat-messages').innerHTML = '<p style="color: #7f8c8d;">Chat messages will appear here...</p>';
    }
    </script>
    """
    display(HTML(html_code))

# Display the chat widget
# create_chat_widget()

def save_conversation(filename="conversation.json"):
    """Save the current conversation to a file"""
    import json
    data = {
        "system_message": chat.system_message,
        "conversation": chat.conversation_history,
        "timestamp": datetime.now().isoformat()
    }
    with open(filename, 'w') as f:
        json.dump(data, f, indent=2)
    display_status(f"✅ Conversation saved to {filename}", "success")

def load_conversation(filename="conversation.json"):
    """Load a conversation from a file"""
    import json
    try:
        with open(filename, 'r') as f:
            data = json.load(f)
        chat.system_message = data["system_message"]
        chat.conversation_history = data["conversation"]
        display_status(f"✅ Conversation loaded from {filename}", "success")
        display_status(f"Loaded {len(data['conversation'])//2} exchanges from {data['timestamp']}", "info")
    except FileNotFoundError:
        display_status(f"❌ File {filename} not found", "error")
    except Exception as e:
        display_status(f"❌ Error loading conversation: {str(e)}", "error")

# Example: Save current conversation
# save_conversation("my_chat.json")

# Example: Load a conversation
# load_conversation("my_chat.json")

# ============================================
# FINAL MESSAGE
# ============================================
display(HTML("""
<div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 30px; border-radius: 15px; margin: 20px 0; text-align: center;">
    <h1>🎉 GPT-OSS:20B is Ready!</h1>
    <p style="font-size: 18px;">Your model is set up and ready to use.</p>
    <p>Use <code>chat.chat("your message")</code> to interact with the model</p>
    <p>Or uncomment <code>chat.interactive_chat()</code> for an interactive session</p>
</div>
"""))

chat.chat("What is the capital of Germany?", include_history=False)

print(chat.conversation_history)

pirate_chat.chat("What is the capital of Germany?", include_history=False)

quick_chat("A list of 10 Oscar Wilde's Books", system = "You are an English Aristocrat of 1800's")

# chat.set_system_message()
chat.clear_history()
chat.set_system_message("You are an helpful assistant. Reasoning: low! ") # Reasoning Levels: High, Medium, Low
chat.chat("Lista di film Marvel dopo Avengers: Endgame")