"""
Bitcoin AI Assistant - Inference Server
Runs the fine-tuned model via llama.cpp with web interface.
"""

from flask import Flask, request, jsonify
from llama_cpp import Llama
import threading
import time

app = Flask(__name__)

# Model configuration
MODEL_PATH = "/app/model/bitcoin-ai-q4.gguf"
N_CTX = 8192  # Context window
N_THREADS = 4  # Match NUC cores

# Global model instance
llm = None
model_ready = False
model_loading = False

def load_model():
    """Load the Bitcoin AI model."""
    global llm, model_ready, model_loading
    model_loading = True
    
    try:
        print(f"Loading model from {MODEL_PATH}...")
        llm = Llama(
            model_path=MODEL_PATH,
            n_ctx=N_CTX,
            n_threads=N_THREADS,
            n_gpu_layers=0,  # CPU only
            verbose=True,
        )
        model_ready = True
        print("Model loaded successfully")
    except Exception as e:
        print(f"Failed to load model: {e}")
        model_ready = False
    finally:
        model_loading = False

@app.route('/health')
def health():
    """Health check endpoint."""
    return jsonify({
        'status': 'healthy',
        'model_ready': model_ready,
        'model_loading': model_loading
    })

@app.route('/chat', methods=['POST'])
def chat():
    """Chat endpoint for Bitcoin AI assistant."""
    if not model_ready:
        return jsonify({
            'error': 'Model not ready',
            'status': model_loading
        }), 503
    
    data = request.json
    prompt = data.get('prompt', '')
    system_prompt = data.get('system', 'You are a Bitcoin operations expert assistant.')
    
    # Build prompt
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    
    # Generate response
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=1024,
        temperature=0.7,
        top_p=0.95,
        repeat_penalty=1.1
    )
    
    return jsonify({
        'response': response['choices'][0]['message']['content'],
        'model': 'bitcoin-ai-7b',
        'tokens_used': response.get('usage', {})
    })

@app.route('/complete', methods=['POST'])
def complete():
    """Completion endpoint for raw text generation."""
    if not model_ready:
        return jsonify({'error': 'Model not ready', 'status': model_loading}), 503
    
    data = request.json
    prompt = data.get('prompt', '')
    max_tokens = data.get('max_tokens', 512)
    
    output = llm(
        prompt,
        max_tokens=max_tokens,
        temperature=0.7,
        top_p=0.95,
        repeat_penalty=1.1,
        echo=False
    )
    
    return jsonify({
        'response': output['choices'][0]['text']
    })

if __name__ == '__main__':
    print("Starting Bitcoin AI Server")
    print("Loading model in background thread...")
    
    # Load model asynchronously
    loader = threading.Thread(target=load_model, daemon=True)
    loader.start()
    
    # Wait for model to load
    while not model_ready and model_loading:
        time.sleep(5)
    
    if model_ready:
        print("Model ready. Starting server...")
        from waitress import serve
        serve(app, host='0.0.0.0', port=8080)
    else:
        print("Model failed to load, exiting")