"""
Entity Extraction Flask Endpoint
Add to airagagent/app.py
"""

from flask import Flask, request, jsonify
from memory.entity_extractor import EntityExtractor
import logging

logger = logging.getLogger(__name__)

# Initialize entity extractor
entity_extractor = EntityExtractor()


@app.route('/extract-entities', methods=['POST'])
def extract_entities_endpoint():
    """
    Extract user entities from conversation
    
    POST body:
    {
        "user_message": "I am a project manager...",
        "assistant_response": "Great! I can help...",
        "use_llm": false  // optional, default false (uses simple patterns)
    }
    
    Returns:
    {
        "success": true,
        "entities": {
            "demographics": {"profession": "project manager"},
            "interests": ["cannabis", "business"],
            "survey_preferences": {},
            "expertise_level": "intermediate"
        },
        "method": "simple_patterns" | "llm",
        "confidence": 0.8
    }
    """
    try:
        data = request.get_json()
        
        if not data or 'user_message' not in data:
            return jsonify({
                'success': False,
                'error': 'user_message required'
            }), 400
        
        user_message = data['user_message']
        assistant_response = data.get('assistant_response', '')
        use_llm = data.get('use_llm', False)
        
        # Extract entities
        if use_llm:
            entities = entity_extractor.extract_entities(user_message, assistant_response)
            method = 'llm'
            confidence = 0.8
        else:
            entities = entity_extractor.extract_simple_patterns(user_message)
            method = 'simple_patterns'
            confidence = 0.6
        
        logger.info(f"Extracted entities via {method}: {entities}")
        
        return jsonify({
            'success': True,
            'entities': entities,
            'method': method,
            'confidence': confidence
        })
        
    except Exception as e:
        logger.error(f"Entity extraction failed: {e}")
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/update-user-profile', methods=['POST'])
def update_user_profile_endpoint():
    """
    Update user profile with extracted entities
    
    POST body:
    {
        "user_id": "user123",
        "entities": {...}
    }
    """
    try:
        data = request.get_json()
        
        if not data or 'user_id' not in data or 'entities' not in data:
            return jsonify({
                'success': False,
                'error': 'user_id and entities required'
            }), 400
        
        # This would integrate with your MySQL user_profiles table
        # For now, just return success
        logger.info(f"Would update profile for {data['user_id']} with {data['entities']}")
        
        return jsonify({
            'success': True,
            'message': 'Profile updated'
        })
        
    except Exception as e:
        logger.error(f"Profile update failed: {e}")
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500