app.py

from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
import uuid
import json
from datetime import datetime
import csv
import io
import numpy as np
import random

from config import Config
from models import db, Session, PageRecord, MouseEvent, FrameCapture
from ml_models.mouse_analyzer import MousePatternAnalyzer
from ml_models.fusion_model import RiskFusionModel
from ml_models.model_loader import model_loader
from ml_models.emotion_engine import EmotionEngine


app = Flask(__name__)
app.config.from_object(Config)
CORS(app)

db.init_app(app)

# تحميل النماذج المدربة
def load_trained_models():
    models_loaded = model_loader.load_all_models(
        emotion_path='ml_models/trained_models/emotion_model.pth',
        mouse_path='ml_models/trained_models/mouse_model.pkl', 
        fusion_path='ml_models/trained_models/fusion_model.pkl'
    )
    
    if models_loaded:
        print("🎉 All trained models loaded successfully!")
    else:
        print("⚠️ Some models failed to load, using fallback methods")

load_trained_models()

emotion_engine = EmotionEngine()
mouse_analyzer = MousePatternAnalyzer()
fusion_model = RiskFusionModel()

# صفحات الاختبار المحدثة (واقعية ومتنوعة)
TEST_PAGES = [
    # صفحات البنوك
    {'id': 'bank_secure', 'type': 'legitimate', 'url': '/pages/bank-secure', 'name': 'البنك الأهلي السعودي'},
    {'id': 'bank_phishing', 'type': 'phishing', 'url': '/pages/bank-phishing', 'name': 'تنبيه البنك السعودي'},
    
    # صفحات البريد الإلكتروني
    {'id': 'email_secure', 'type': 'legitimate', 'url': '/pages/email-secure', 'name': 'Outlook - البريد'},
    {'id': 'email_phishing', 'type': 'phishing', 'url': '/pages/email-phishing', 'name': 'رسالة خدمة العملاء'},
    
    # صفحات وسائل التواصل
    {'id': 'social_secure', 'type': 'legitimate', 'url': '/pages/social-secure', 'name': 'Facebook تسجيل الدخول'},
    {'id': 'social_phishing', 'type': 'phishing', 'url': '/pages/social-phishing', 'name': 'تنبيه فيسبوك الأمني'},
    
    # صفحات التسوق
    {'id': 'shopping_secure', 'type': 'legitimate', 'url': '/pages/shopping-secure', 'name': 'Amazon تسجيل الدخول'},
    {'id': 'shopping_phishing', 'type': 'phishing', 'url': '/pages/shopping-phishing', 'name': 'عرض خاص - خصم 80%'},
    
    # صفحات الخدمات الحكومية
    {'id': 'gov_secure', 'type': 'legitimate', 'url': '/pages/gov-secure', 'name': 'أبشر - الخدمات الإلكترونية'},
    {'id': 'gov_phishing', 'type': 'phishing', 'url': '/pages/gov-phishing', 'name': 'تنبيه وزارة الداخلية'}
]

@app.route('/')
def home():
    return jsonify({
        'message': 'Phishing Study Backend API',
        'status': 'running',
        'endpoints': {
            'health': '/api/health',
            'start_session': '/api/session/start (POST)',
            'submit_page': '/api/session/<session_id>/page (POST)',
            'end_session': '/api/session/<session_id>/end (POST)',
            'export_data': '/api/admin/export (GET)'
        }
    })

@app.route('/api/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy', 'timestamp': datetime.utcnow().isoformat()})

@app.route('/api/session/start', methods=['POST'])
def start_session():
    try:
        data = request.get_json() or {}
        user_id = data.get('user_id', f'user_{uuid.uuid4().hex[:8]}')
        username = data.get('username')
        session_id = f'sess_{uuid.uuid4().hex[:16]}'
        pages_order = random.sample(TEST_PAGES, len(TEST_PAGES))
        
        session = Session(
            id=session_id,
            user_id=user_id,
            username=username,
            consent_given=True,
            pages_order=json.dumps(pages_order)
        )
        
        db.session.add(session)
        db.session.commit()
        
        return jsonify({
            'session_id': session_id,
            'user_id': user_id,
            'username': username,
            'pages': pages_order,
            'start_time': session.start_time.isoformat()
        }), 201
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/session/<session_id>/page', methods=['POST'])
def submit_page_data(session_id):
    try:
        data = request.get_json()
        if not data:
            return jsonify({'error': 'No data provided'}), 400

        # ===== Session =====
        session = Session.query.get(session_id)
        if not session:
            return jsonify({'error': 'Session not found'}), 404

        # ===== Times =====
        start_time = datetime.fromisoformat(data['start_time'].replace('Z', '+00:00'))
        end_time = datetime.fromisoformat(data['end_time'].replace('Z', '+00:00'))

        # ===== Frames =====
        frames = data.get('frames', [])
        if not frames:
            return jsonify({'error': 'No frames received'}), 400

        # ناخد آخر فريم بس (أنضف وأسرع)
        last_frame = frames[-1]
        base64_img = last_frame.get('img_base64')

        if not base64_img:
            return jsonify({'error': 'No image data in frame'}), 400

        # ===== Emotion (PyTorch ViT Engine) =====
        emotion_probs, dominant_emotion = emotion_engine.predict_from_base64(
            base64_img
        )

        # حساب الـ emotion risk
        emotion_risk = (
            emotion_probs.get('fear', 0) * 0.6 +
            emotion_probs.get('sad', 0) * 0.3 +
            emotion_probs.get('angry', 0) * 0.1
        )
        emotion_risk = min(emotion_risk, 1.0)

        # Debug (مفيد في التيرمنال)
        print("======== EMOTION ENGINE DEBUG ========")
        print("Dominant emotion:", dominant_emotion)
        print("Emotion probabilities:")
        for k, v in emotion_probs.items():
            print(f"  {k}: {v:.4f}")
        print("Emotion risk:", emotion_risk)
        print("=====================================")

        # ===== Mouse =====
        mouse_events = data.get('mouse_events', [])
        mouse_features = mouse_analyzer.extract_features(
            mouse_events, start_time, end_time
        )
        mouse_risk = mouse_analyzer.calculate_mouse_risk(mouse_features)

        # ===== Fusion =====
        phishing_score = fusion_model.calculate_phishing_score(
            emotion_risk, mouse_risk, mouse_features
        )

        # ===== PageRecord =====
        page_record = PageRecord(
            id=f'page_rec_{uuid.uuid4().hex[:16]}',
            session_id=session_id,
            page_id=data.get('page_id'),
            page_type=data.get('page_type'),
            start_time=start_time,
            end_time=end_time,
            user_label=data.get('label'),
            notes=data.get('notes', '')
        )

        page_record.emotion_probs = json.dumps(emotion_probs)
        page_record.dominant_emotion = dominant_emotion
        page_record.emotion_risk = emotion_risk
        page_record.mouse_features = json.dumps(mouse_features)
        page_record.mouse_risk = mouse_risk
        page_record.phishing_score = phishing_score

        db.session.add(page_record)

        # ===== Mouse Events =====
        for event in mouse_events:
            mouse_event = MouseEvent(
                page_record_id=page_record.id,
                event_type=event['type'],
                x=event.get('x'),
                y=event.get('y'),
                timestamp=datetime.fromisoformat(event['t'].replace('Z', '+00:00')),
                additional_data=json.dumps({
                    k: v for k, v in event.items()
                    if k not in ['type', 'x', 'y', 't']
                })
            )
            db.session.add(mouse_event)

        db.session.commit()

        # ===== Response =====
        return jsonify({
            'emotion_source': 'internal_engine',
            'emotion_probs': emotion_probs,
            'dominant_emotion': dominant_emotion,
            'emotion_risk': emotion_risk,
            'mouse_features': mouse_features,
            'mouse_risk': mouse_risk,
            'phishing_score': phishing_score,
            'risk_level': fusion_model.get_risk_level(phishing_score)
        }), 201

    except Exception as e:
        db.session.rollback()
        return jsonify({'error': str(e)}), 500


@app.route('/api/admin/export', methods=['GET'])
def export_data():
    try:
        records = PageRecord.query.all()
        
        output = io.StringIO()
        writer = csv.writer(output)
        
        writer.writerow([
            'record_id', 'user_id', 'username', 'session_id', 'page_id', 'page_type',
            'start_time', 'timestamp', 'sample_type', 'emotion_probs',
            'mouse_features', 'emotion_risk', 'mouse_risk', 'phishing_score',
            'label', 'dominant_emotion', 'notes'
        ])
        
        for record in records:
            writer.writerow([
                record.id,
                record.session.user_id,
                record.session.username,
                record.session_id,
                record.page_id,
                record.page_type,
                record.start_time.isoformat() if record.start_time else '',
                record.timestamp.isoformat() if record.timestamp else '',
                'detail',
                record.emotion_probs or '{}',
                record.mouse_features or '{}',
                record.emotion_risk or 0.0,
                record.mouse_risk or 0.0,
                record.phishing_score or 0.0,
                record.user_label or '',
                record.dominant_emotion or '',
                record.notes or ''
            ])
        
        output.seek(0)
        return send_file(
            io.BytesIO(output.getvalue().encode('utf-8-sig')),
            mimetype='text/csv; charset=utf-8',
            as_attachment=True,
            download_name=f'phishing_study_export_{datetime.utcnow().strftime("%Y%m%d_%H%M%S")}.csv'
        )    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/session/<session_id>', methods=['GET'])
def get_session(session_id):
    session = Session.query.get(session_id)
    if not session:
        return jsonify({'error': 'Session not found'}), 404
    
    page_records = PageRecord.query.filter_by(session_id=session_id).all()
    
    return jsonify({
        'session_id': session.id,
        'user_id': session.user_id,
        'username': session.username,
        'start_time': session.start_time.isoformat(),
        'end_time': session.end_time.isoformat() if session.end_time else None,
        'completed': session.completed,
        'page_records': [
            {
                'page_id': pr.page_id,
                'user_label': pr.user_label,
                'phishing_score': pr.phishing_score,
                'dominant_emotion': pr.dominant_emotion
            } for pr in page_records
        ]
    })

if __name__ == '__main__':
    with app.app_context():
        db.create_all()
    print("Phishing Study Backend starting on http://localhost:5000")
    app.run(debug=True, host='0.0.0.0', port=5000)