from flask import Flask, render_template, request, jsonify
import os
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import threading

app = Flask(__name__)

# Global variable to store the model
llm = None

def initialize_model():
    global llm
    try:
        print("Initializing TinyLlama model...")
        
        # Model details
        model_name = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
        model_file = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
        
        # Download model if not already present
        print(f"Downloading model: {model_name}/{model_file}")
        model_path = hf_hub_download(model_name, filename=model_file)
        
        print(f"Model downloaded to: {model_path}")
        
        # Initialize the model
        llm = Llama(
            model_path=model_path,
            n_ctx=16000,  # Context length to use
            n_threads=4,  # Number of CPU threads to use
            n_gpu_layers=-1,  # Number of model layers to offload to GPU (-1 for all layers if GPU is available)
            verbose=False  # Disable verbose output
        )
        
        print("Model initialized successfully!")
        
    except Exception as e:
        print(f"Error initializing model: {str(e)}")
        raise

# Initialize model in a separate thread to not block app startup
model_init_thread = threading.Thread(target=initialize_model)
model_init_thread.start()

@app.route('/')
def index():
    # Wait for model to be initialized before serving the page
    if llm is None:
        return "Model is still loading, please wait...", 503
    return render_template('index.html')

@app.route('/health')
def health():
    status = 'healthy' if llm is not None else 'loading'
    return {'status': status}

@app.route('/chat', methods=['POST'])
def chat():
    global llm
    
    # Wait for model to be initialized
    if llm is None:
        if model_init_thread.is_alive():
            return jsonify({'error': 'Model is still loading, please wait...'}), 503
        else:
            return jsonify({'error': 'Model failed to load'}), 500
    
    try:
        user_message = request.json.get('message', '')
        
        if not user_message:
            return jsonify({'error': 'No message provided'}), 400
        
        # Create chat completion
        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": user_message}
            ],
            max_tokens=100,
            temperature=0.7,
            stop=["</s>", "\n"]
        )
        
        # Extract the response text
        assistant_response = response["choices"][0]["message"]["content"]
        
        return jsonify({
            'response': assistant_response.strip()
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8501, debug=True)