import os
import requests
import subprocess
import json
import time
import spaces
import gradio as gr
import random
from typing import List, Optional, Tuple, Dict

# Constants for the llama.cpp server
API_PATH_HEALTH = "/health"
API_PATH_COMPLETIONS = "/chat/completions"
LLAMA_CPP_SERVER_BASE = "http://127.0.0.1:8080"
LLAMA_CPP_SERVER_START_TIMEOUT = 50  # seconds
MODEL_FILENAME = "AstroSage-8B-Q8_0.gguf"
HF_MODEL_ID = "AstroMLab/AstroSage-8B-GGUF"

# Ensure the model is available
if not os.path.exists(MODEL_FILENAME):
    url = f"https://huggingface.co/{HF_MODEL_ID}/resolve/main/{MODEL_FILENAME}"
    subprocess.check_call(["curl", "-o", MODEL_FILENAME, "-L", url])

if not os.path.exists("llama-server"):
    subprocess.check_call("curl -o llama-server -L https://ngxson-llamacpp-builder.hf.space/llama-server", shell=True)
    subprocess.check_call("chmod +x llama-server", shell=True)

# Roles and History Types
class Role:
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"

History = List[Dict[str, str]]  # Chat history with "role" and "content"

# Placeholder greeting messages
GREETING_MESSAGES = [
    "Greetings! I am AstroSage, your guide to the cosmos. What would you like to explore today?",
    "Welcome to our cosmic journey! I am AstroSage. How may I assist you in understanding the universe?",
    "AstroSage here. Ready to explore the mysteries of space and time. How may I be of assistance?",
    "The universe awaits! I'm AstroSage. What astronomical wonders shall we discuss?",
]

# Helper functions
def wait_until_llamacpp_ready():
    """Wait until the llama.cpp server is ready."""
    trials = 0
    while trials < LLAMA_CPP_SERVER_START_TIMEOUT:
        try:
            response = requests.get(LLAMA_CPP_SERVER_BASE + API_PATH_HEALTH)
            if response.status_code == 200:
                return
        except requests.exceptions.RequestException:
            pass
        time.sleep(1)
        trials += 1
    raise TimeoutError("llama.cpp server did not start in time.")

def initial_greeting() -> History:
    """Generate the initial greeting from the assistant."""
    return [{"role": "assistant", "content": random.choice(GREETING_MESSAGES)}]

def send_request_to_llama(query: str, history: History) -> str:
    """Send a chat request to the llama.cpp server."""
    messages = [{"role": Role.SYSTEM, "content": "You are AstroSage, an AI assistant specializing in astronomy, astrophysics, and cosmology."}]
    messages.extend(history)
    messages.append({"role": Role.USER, "content": query})

    headers = {"Content-Type": "application/json"}
    data = {"temperature": 0.7, "messages": messages, "stream": True}
    response = requests.post(LLAMA_CPP_SERVER_BASE + API_PATH_COMPLETIONS, headers=headers, json=data, stream=True)
    response.raise_for_status()

    response_text = ""
    for line in response.iter_lines():
        line = line.decode("utf-8")
        if line.startswith("data: ") and not line.endswith("[DONE]"):
            data = json.loads(line[len("data: "):])
            response_text += data["choices"][0]["delta"].get("content", "")
    return response_text

@spaces.GPU
def bot(history: Optional[History]) -> History:
    """Generate the assistant's response."""
    if history is None:
        history = []
    query = history[-1]["content"]
    
    response = send_request_to_llama(query, history[:-1])
    history.append({"role": "assistant", "content": response})
    return history

# Custom CSS for a space theme
custom_css = """
#component-0 {
    background-color: #1a1a2e;
    border-radius: 15px;
    padding: 20px;
}
.dark {
    background-color: #0f0f1a;
}
.contain {
    max-width: 1200px !important;
}
"""

# Launch llama.cpp server
llama_proc = subprocess.Popen([
    "./llama-server"
], env=dict(
    os.environ,
    LLAMA_HOST="0.0.0.0",
    LLAMA_PORT="8080",
    LLAMA_ARG_CTX_SIZE=str(2048),
    LLAMA_ARG_MODEL=MODEL_FILENAME,
    LLAMA_ARG_N_GPU_LAYERS="9999",
))

try:
    wait_until_llamacpp_ready()

    # Create the Gradio interface
    with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")) as demo:
        gr.Markdown(
            """
            # 🌌 AstroSage: Your Cosmic AI Companion
            
            Welcome to AstroSage, an advanced AI assistant specializing in astronomy, astrophysics, and cosmology. 
            Powered by the AstroSage-Llama-3.1-8B model, I'm here to help you explore the wonders of the universe!
            
            ### What Can I Help You With?
            - 🪐 Explanations of astronomical phenomena
            - 🚀 Space exploration and missions
            - ⭐ Stars, galaxies, and cosmology
            - 🌍 Planetary science and exoplanets
            - 📊 Astrophysics concepts and theories
            - 🔭 Astronomical instruments and observations
            
            Just type your question below and let's embark on a cosmic journey together!
            """
        )
        
        chatbot = gr.Chatbot(
            label="Chat with AstroSage",
            bubble_full_width=False,
            show_label=True,
            height=450,
            type="messages"
        )
        
        with gr.Row():
            msg = gr.Textbox(
                label="Type your message here",
                placeholder="Ask me anything about space and astronomy...",
                scale=9
            )
            clear = gr.Button("Clear Chat", scale=1)
        
        # Example questions for quick start
        gr.Examples(
            examples=[
                "What is a black hole and how does it form?",
                "Can you explain the life cycle of a star?",
                "What are exoplanets and how do we detect them?",
                "Tell me about the James Webb Space Telescope.",
                "What is dark matter and why is it important?"
            ],
            inputs=msg,
            label="Example Questions"
        )
        
        # Set up the message chain
        msg.submit(
            lambda x, y: (x, y + [{"role": "user", "content": x}]),
            [msg, chatbot],
            [msg, chatbot],
            queue=False
        ).then(
            bot,
            chatbot,
            chatbot
        )
        
        # Clear button functionality
        clear.click(lambda: None, None, chatbot, queue=False)
        
        # Initial greeting
        demo.load(initial_greeting, None, chatbot, queue=False)

    # Launch the app
    demo.launch()

finally:
    llama_proc.kill()