import os
import time
import logging
from flask import Flask, request, jsonify, render_template
from datetime import datetime
from src.retrieval import DocumentRetriever
from src.generation import ResponseGenerator
from src.history import ChatHistory

app = Flask(__name__)

# Configuration
CONFIG = {
    'cache_dir': '/cache',
    'data_dir': '/app/data',
    'max_history': 10
}

# Initialize logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Global state
models_loaded = False
start_time = time.time()

def initialize_services():
    global models_loaded
    
    try:
        logger.info("Initializing services...")
        
        # Initialize components
        app.retriever = DocumentRetriever(
            model_name='all-MiniLM-L6-v2',
            cache_folder=CONFIG['cache_dir'],
            data_path=os.path.join(CONFIG['data_dir'], 'rupeia_document.json')
        )
        
        app.generator = ResponseGenerator(
            model_name='distilgpt2',
            cache_dir=CONFIG['cache_dir']
        )
        
        app.history = ChatHistory(
            storage_path=os.path.join(CONFIG['data_dir'], 'chat_history.json'),
            max_entries=CONFIG['max_history']
        )
        
        models_loaded = True
        logger.info(f"Services initialized in {time.time()-start_time:.2f}s")
    
    except Exception as e:
        logger.error(f"Initialization failed: {str(e)}")
        raise

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/ask', methods=['POST'])
def chat():
    if not models_loaded:
        return jsonify({
            "status": "initializing",
            "message": "System is starting up, please try again shortly"
        }), 503
    
    try:
        user_input = request.json.get('message', '').strip()
        if not user_input:
            return jsonify({"error": "Empty message"}), 400
        
        # Retrieve relevant context
        context = app.retriever.retrieve(user_input)
        
        # Generate response
        response = app.generator.generate(
            prompt=user_input,
            context=context
        )
        
        # Save to history
        app.history.add_entry(
            user_message=user_input,
            bot_response=response
        )
        
        return jsonify({
            "response": response,
            "timestamp": datetime.utcnow().isoformat(),
            "quickReplies": []  # You can add logic here to provide quick replies if needed
        })
    
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        return jsonify({
            "error": "An error occurred",
            "details": str(e)
        }), 500

@app.route('/api/history', methods=['GET'])
def get_history():
    return jsonify(app.history.get_history())

if __name__ == '__main__':
    initialize_services()  # Initialize services before starting the app
    app.run(host='0.0.0.0', port=8000)