"""
Main Flask application for the Malware Website Crawler
"""
import os
import logging
from flask import Flask, render_template, request, jsonify, redirect, url_for
from flask_cors import CORS
from celery import Celery
from database import init_db, db_session
from models import ScanJob, MalwareFinding, ContextWord
from crawler import WebsiteCrawler
from analyzer import MalwareAnalyzer
from config import Config

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def create_app():
    app = Flask(__name__)
    app.config.from_object(Config)
    CORS(app)
    
    # Initialize database
    init_db()
    
    # Initialize Celery
    celery = make_celery(app)
    
    @app.route('/')
    def index():
        return render_template('index.html')
    
    @app.route('/api/scan', methods=['POST'])
    def start_scan():
        """Start a new website scan"""
        data = request.get_json()
        target_url = data.get('url')
        scan_type = data.get('type', 'website')
        
        if not target_url:
            return jsonify({'error': 'URL is required'}), 400
        
        # Create scan job
        job = ScanJob(
            target_url=target_url,
            scan_type=scan_type,
            status='pending'
        )
        db_session.add(job)
        db_session.commit()
        
        # Start background task
        task = crawl_and_analyze.delay(job.id)
        job.task_id = task.id
        db_session.commit()
        
        return jsonify({
            'job_id': job.id,
            'task_id': task.id,
            'status': 'started'
        })
    
    @app.route('/api/scan/<int:job_id>')
    def get_scan_status(job_id):
        """Get scan status and results"""
        job = db_session.query(ScanJob).filter_by(id=job_id).first()
        if not job:
            return jsonify({'error': 'Job not found'}), 404
        
        findings = db_session.query(MalwareFinding).filter_by(scan_job_id=job_id).all()
        
        return jsonify({
            'job': {
                'id': job.id,
                'target_url': job.target_url,
                'status': job.status,
                'created_at': job.created_at.isoformat(),
                'completed_at': job.completed_at.isoformat() if job.completed_at else None
            },
            'findings': [{
                'id': f.id,
                'filename': f.filename,
                'file_hash': f.file_hash,
                'malware_name': f.malware_name,
                'confidence': f.confidence,
                'features': f.features,
                'ip_addresses': f.ip_addresses,
                'attack_vectors': f.attack_vectors,
                'context_words': f.context_words
            } for f in findings]
        })
    
    @app.route('/api/reports')
    def get_reports():
        """Get all scan reports"""
        jobs = db_session.query(ScanJob).order_by(ScanJob.created_at.desc()).all()
        return jsonify([{
            'id': job.id,
            'target_url': job.target_url,
            'status': job.status,
            'created_at': job.created_at.isoformat(),
            'findings_count': len(job.findings)
        } for job in jobs])
    
    @app.route('/reports')
    def reports():
        return render_template('reports.html')
    
    return app

def make_celery(app):
    """Initialize Celery with Redis backend"""
    celery = Celery(
        app.import_name,
        backend=app.config['CELERY_RESULT_BACKEND'],
        broker=app.config['CELERY_BROKER_URL']
    )
    
    class ContextTask(celery.Task):
        def __call__(self, *args, **kwargs):
            with app.app_context():
                return self.run(*args, **kwargs)
    
    celery.Task = ContextTask
    return celery

# Create Celery instance
app = create_app()
celery = make_celery(app)

@celery.task(bind=True)
def crawl_and_analyze(self, job_id):
    """Background task to crawl website and analyze files"""
    try:
        job = db_session.query(ScanJob).filter_by(id=job_id).first()
        if not job:
            return {'error': 'Job not found'}
        
        job.status = 'running'
        db_session.commit()
        
        # Initialize crawler and analyzer
        crawler = WebsiteCrawler()
        analyzer = MalwareAnalyzer()
        
        # Crawl website for files
        files = crawler.crawl_website(job.target_url)
        
        findings = []
        for file_info in files:
            # Update progress
            self.update_state(
                state='PROGRESS',
                meta={'current': len(findings), 'total': len(files)}
            )
            
            # Analyze file
            analysis = analyzer.analyze_file(file_info)
            if analysis:
                finding = MalwareFinding(
                    scan_job_id=job_id,
                    filename=file_info['filename'],
                    file_path=file_info['path'],
                    file_hash=analysis['hash'],
                    malware_name=analysis.get('name'),
                    confidence=analysis.get('confidence', 0),
                    features=analysis.get('features', []),
                    ip_addresses=analysis.get('ip_addresses', []),
                    attack_vectors=analysis.get('attack_vectors', []),
                    context_words=analysis.get('context_words', [])
                )
                db_session.add(finding)
                findings.append(finding)
        
        job.status = 'completed'
        job.completed_at = db_session.query(db_session.func.now()).scalar()
        db_session.commit()
        
        return {'status': 'completed', 'findings_count': len(findings)}
        
    except Exception as e:
        logger.error(f"Error in crawl_and_analyze: {str(e)}")
        job.status = 'failed'
        db_session.commit()
        raise

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
