"""
Lab Report Decoder - Flask Application
Professional web interface for lab report analysis
Fixed for Hugging Face Spaces stateless environment
"""

from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
import os
import tempfile
import secrets
import json
import uuid
from pdf_extractor import LabReportExtractor, LabResult
from rag_engine import LabReportRAG
from dotenv import load_dotenv

load_dotenv()

app = Flask(__name__)
app.secret_key = os.getenv('SECRET_KEY', secrets.token_hex(16))
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
app.config['UPLOAD_FOLDER'] = tempfile.gettempdir()

# Initialize RAG system (singleton)
rag_system = None

# In-memory storage for sessions (better than Flask sessions in HF Spaces)
session_storage = {}

def get_rag_system():
    """Lazy load RAG system"""
    global rag_system
    if rag_system is None:
        print("🔄 Initializing RAG system...")
        rag_system = LabReportRAG()
        print("✅ RAG system ready")
    return rag_system

@app.route('/')
def index():
    """Main page"""
    return render_template('index.html')

@app.route('/api/upload', methods=['POST'])
def upload_file():
    """Handle PDF upload and extraction"""
    try:
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        if not file.filename.lower().endswith('.pdf'):
            return jsonify({'error': 'Only PDF files are allowed'}), 400
        
        # Save file temporarily
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        try:
            # Extract lab results
            print("📄 Extracting lab results from PDF...")
            extractor = LabReportExtractor()
            results = extractor.extract_from_pdf(filepath)
            
            if not results:
                return jsonify({'error': 'No lab results found in PDF. Please make sure your PDF contains a valid lab report with test names, values, and reference ranges.'}), 400
            
            print(f"✅ Extracted {len(results)} results")
            
            # Convert to JSON-serializable format
            results_data = [
                {
                    'test_name': r.test_name,
                    'value': r.value,
                    'unit': r.unit,
                    'reference_range': r.reference_range,
                    'status': r.status
                }
                for r in results
            ]
            
            # Generate a unique session ID
            session_id = str(uuid.uuid4())
            
            # Store results in memory with session ID
            session_storage[session_id] = {
                'results': results_data,
                'results_objects': results  # Store LabResult objects for later use
            }
            
            print(f"💾 Stored results in session: {session_id}")
            
            return jsonify({
                'success': True,
                'session_id': session_id,
                'results': results_data,
                'count': len(results_data)
            })
        
        finally:
            # Clean up temp file
            if os.path.exists(filepath):
                os.remove(filepath)
    
    except Exception as e:
        print(f"❌ Upload error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/explain', methods=['POST'])
def explain_results():
    """Generate explanations for lab results"""
    try:
        data = request.get_json()
        session_id = data.get('session_id')
        
        if not session_id or session_id not in session_storage:
            return jsonify({'error': 'Session expired. Please upload your PDF again.'}), 400
        
        # Get results from storage
        session_data = session_storage[session_id]
        results = session_data['results_objects']
        
        print(f"🧠 Generating explanations for {len(results)} results...")
        
        # Generate explanations
        rag = get_rag_system()
        explanations = {}
        
        for i, result in enumerate(results):
            print(f"  Explaining {i+1}/{len(results)}: {result.test_name}...")
            try:
                explanation = rag.explain_result(result)
                explanations[result.test_name] = explanation
            except Exception as e:
                print(f"    Error: {str(e)}")
                explanations[result.test_name] = f"Unable to generate explanation: {str(e)}"
        
        # Store explanations in session
        session_storage[session_id]['explanations'] = explanations
        
        print("✅ All explanations generated")
        
        return jsonify({
            'success': True,
            'explanations': explanations
        })
    
    except Exception as e:
        print(f"❌ Explanation error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/ask', methods=['POST'])
def ask_question():
    """Answer follow-up questions"""
    try:
        data = request.get_json()
        question = data.get('question', '').strip()
        session_id = data.get('session_id')
        
        if not question:
            return jsonify({'error': 'No question provided'}), 400
        
        if not session_id or session_id not in session_storage:
            return jsonify({'error': 'Session expired. Please upload your PDF again.'}), 400
        
        # Get results from storage
        session_data = session_storage[session_id]
        results = session_data['results_objects']
        
        print(f"💬 Answering question: {question}")
        
        # Get answer
        rag = get_rag_system()
        answer = rag.answer_followup_question(question, results)
        
        print("✅ Answer generated")
        
        return jsonify({
            'success': True,
            'question': question,
            'answer': answer
        })
    
    except Exception as e:
        print(f"❌ Question error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/summary', methods=['POST'])
def get_summary():
    """Generate overall summary"""
    try:
        data = request.get_json()
        session_id = data.get('session_id')
        
        if not session_id or session_id not in session_storage:
            return jsonify({'error': 'Session expired. Please upload your PDF again.'}), 400
        
        # Get results from storage
        session_data = session_storage[session_id]
        results = session_data['results_objects']
        
        print("📊 Generating summary...")
        
        # Generate summary
        rag = get_rag_system()
        summary = rag.generate_summary(results)
        
        # Calculate statistics
        stats = {
            'total': len(results),
            'normal': sum(1 for r in results if r.status == 'normal'),
            'high': sum(1 for r in results if r.status == 'high'),
            'low': sum(1 for r in results if r.status == 'low'),
            'unknown': sum(1 for r in results if r.status == 'unknown')
        }
        
        print(f"✅ Summary generated - Stats: {stats}")
        
        return jsonify({
            'success': True,
            'summary': summary,
            'stats': stats
        })
    
    except Exception as e:
        print(f"❌ Summary error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'active_sessions': len(session_storage)
    })

@app.errorhandler(413)
def request_entity_too_large(error):
    return jsonify({'error': 'File too large. Maximum size is 16MB.'}), 413

@app.errorhandler(500)
def internal_error(error):
    return jsonify({'error': 'Internal server error'}), 500

if __name__ == '__main__':
    print("🚀 Starting Lab Report Decoder...")
    print("📍 Server will be available at http://0.0.0.0:7860")
    app.run(debug=True, host='0.0.0.0', port=7860)