from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import os
import time
import traceback
from pathlib import Path
import threading
import atexit
import random
import math
from PIL import Image, ImageDraw
import requests
import uuid
from datetime import datetime, timedelta
from collections import deque
import string

app = Flask(__name__)
CORS(app)

BASE_DIR = Path(__file__).parent
PUBLIC_DIR = BASE_DIR / 'public'
PUBLIC_DIR.mkdir(exist_ok=True)

PORT = int(os.environ.get('PORT', 7860))

VIEWPORT_WIDTH = 1920
VIEWPORT_HEIGHT = 1080

VIEWPORT_CONFIG = {
    'width': VIEWPORT_WIDTH,
    'height': VIEWPORT_HEIGHT,
    'device_scale_factor': 1
}

MAX_ROOMS = 5
ROOM_TIMEOUT_MINUTES = 10
SCREENSHOT_EXPIRY_MINUTES = 10
JOB_EXPIRY_MINUTES = 30

SCREENSHOT_BASE_URL = "http://pnode1.danbot.host:1149"

class Job:
    def __init__(self, job_id, code, lang, base_url):
        self.job_id = job_id
        self.code = code
        self.lang = lang
        self.base_url = base_url
        self.status = 'queued'
        self.result = None
        self.error = None
        self.created_at = datetime.now()
        self.started_at = None
        self.completed_at = None
        self.room_id = None
        self.screenshots = []
        
    def to_dict(self):
        return {
            'job_id': self.job_id,
            'status': self.status,
            'result': self.result,
            'error': self.error,
            'created_at': self.created_at.isoformat(),
            'started_at': self.started_at.isoformat() if self.started_at else None,
            'completed_at': self.completed_at.isoformat() if self.completed_at else None,
            'room_id': self.room_id,
            'screenshots': self.screenshots
        }
        
    def is_expired(self):
        return (datetime.now() - self.created_at) > timedelta(minutes=JOB_EXPIRY_MINUTES)

class BrowserRoom:
    def __init__(self, room_id):
        self.room_id = room_id
        self.browser = None
        self.context = None
        self.page = None
        self.cookies = []
        self.created_at = datetime.now()
        self.last_activity = datetime.now()
        self.is_busy = False
        self.current_job_id = None
        self.lock = threading.Lock()
        self.screenshots_dir = PUBLIC_DIR / f'room_{room_id}'
        self.screenshots_dir.mkdir(exist_ok=True)
        
    def update_activity(self):
        self.last_activity = datetime.now()
        
    def is_expired(self):
        return (datetime.now() - self.last_activity) > timedelta(minutes=ROOM_TIMEOUT_MINUTES)
        
    def cleanup_screenshots(self):
        try:
            current_time = time.time()
            for file in self.screenshots_dir.glob('*.png'):
                file_age = current_time - file.stat().st_mtime
                if file_age > (SCREENSHOT_EXPIRY_MINUTES * 60):
                    file.unlink()
                    print(f"[ROOM-{self.room_id}] Deleted expired screenshot: {file.name}")
        except Exception as e:
            print(f"[ROOM-{self.room_id}] Screenshot cleanup error: {e}")
    
    def reset_browser(self):
        try:
            if self.page:
                self.page.close()
                self.page = None
            if self.context:
                try:
                    self.cookies = self.context.cookies()
                    print(f"[ROOM-{self.room_id}] Saved {len(self.cookies)} cookies")
                except:
                    pass
                self.context.close()
                self.context = None
            if self.browser:
                self.browser.close()
                self.browser = None
        except Exception as e:
            print(f"[ROOM-{self.room_id}] Browser reset error: {e}")
            
    def cleanup(self):
        try:
            self.reset_browser()
            self.cleanup_screenshots()
            if self.screenshots_dir.exists():
                try:
                    self.screenshots_dir.rmdir()
                except:
                    pass
        except Exception as e:
            print(f"[ROOM-{self.room_id}] Cleanup error: {e}")

class JobManager:
    def __init__(self):
        self.jobs = {}
        self.lock = threading.Lock()
        
    def generate_job_id(self):
        while True:
            job_id = ''.join(random.choices(string.ascii_uppercase + string.digits, k=3))
            if job_id not in self.jobs:
                return job_id
                
    def create_job(self, code, lang, base_url):
        with self.lock:
            job_id = self.generate_job_id()
            job = Job(job_id, code, lang, base_url)
            self.jobs[job_id] = job
            return job
            
    def get_job(self, job_id):
        with self.lock:
            return self.jobs.get(job_id)
            
    def update_job(self, job_id, **kwargs):
        with self.lock:
            if job_id in self.jobs:
                job = self.jobs[job_id]
                for key, value in kwargs.items():
                    setattr(job, key, value)
                    
    def cleanup_expired_jobs(self):
        with self.lock:
            expired = [jid for jid, job in self.jobs.items() if job.is_expired()]
            for jid in expired:
                del self.jobs[jid]
                print(f"[JOB-{jid}] Deleted expired job")

class RoomManager:
    def __init__(self):
        self.rooms = {}
        self.available_rooms = deque(range(1, MAX_ROOMS + 1))
        self.lock = threading.Lock()
        self.cleanup_thread = None
        self.stop_cleanup = threading.Event()
        
    def acquire_room(self, job_id, timeout=300):
        start_time = time.time()
        while True:
            with self.lock:
                if self.available_rooms:
                    room_id = self.available_rooms.popleft()
                    if room_id not in self.rooms:
                        self.rooms[room_id] = BrowserRoom(room_id)
                    room = self.rooms[room_id]
                    room.is_busy = True
                    room.current_job_id = job_id
                    room.update_activity()
                    print(f"[ROOM-{room_id}] Acquired for job {job_id}. Available rooms: {len(self.available_rooms)}/{MAX_ROOMS}")
                    return room
                    
            if time.time() - start_time > timeout:
                raise Exception(f"Timeout waiting for available room after {timeout}s. All {MAX_ROOMS} rooms are busy.")
                
            print(f"[QUEUE] All {MAX_ROOMS} rooms busy. Waiting for available room...")
            time.sleep(1)
            
    def release_room(self, room):
        with self.lock:
            room.is_busy = False
            room.current_job_id = None
            room.update_activity()
            if room.room_id not in self.available_rooms:
                self.available_rooms.append(room.room_id)
            print(f"[ROOM-{room.room_id}] Released. Available rooms: {len(self.available_rooms)}/{MAX_ROOMS}")
            
    def cleanup_expired_rooms(self):
        while not self.stop_cleanup.is_set():
            try:
                current_time = time.time()
                
                with self.lock:
                    for room_id, room in list(self.rooms.items()):
                        if not room.is_busy:
                            try:
                                for file in room.screenshots_dir.glob('*.png'):
                                    file_age = current_time - file.stat().st_mtime
                                    if file_age > (SCREENSHOT_EXPIRY_MINUTES * 60):
                                        file.unlink()
                                        print(f"[ROOM-{room_id}] Auto-deleted expired screenshot: {file.name}")
                            except Exception as e:
                                print(f"[ROOM-{room_id}] Auto cleanup screenshot error: {e}")
                        
                        if not room.is_busy and room.is_expired():
                            print(f"[ROOM-{room_id}] Cleaning up expired room")
                            self.rooms[room_id].cleanup()
                            del self.rooms[room_id]
                        
            except Exception as e:
                print(f"[CLEANUP] Error: {e}")
            self.stop_cleanup.wait(60)
            
    def start_cleanup(self):
        self.cleanup_thread = threading.Thread(target=self.cleanup_expired_rooms, daemon=True)
        self.cleanup_thread.start()
        
    def stop_cleanup_thread(self):
        self.stop_cleanup.set()
        if self.cleanup_thread:
            self.cleanup_thread.join(timeout=5)
            
    def shutdown_all(self):
        with self.lock:
            for room in self.rooms.values():
                room.cleanup()
            self.rooms.clear()
            self.available_rooms.clear()

room_manager = RoomManager()
job_manager = JobManager()

def execute_in_room(code_snippet, room, job_id):
    result = {'screenshots': [], 'data': None, 'error': None}
    
    room.reset_browser()
    
    browser_obj = None
    
    try:
        print(f"[ROOM-{room.room_id}] [JOB-{job_id}] Starting execution at {time.time()}")
        
        from camoufox.sync_api import Camoufox
        
        proxy_server = "http://pnode1.danbot.host:1271"
        
        browser_obj = Camoufox(
            headless=True,
            humanize=True,
            proxy={'server': proxy_server}
        )
        
        browser_obj = browser_obj.__enter__()
        
        if room.cookies:
            print(f"[ROOM-{room.room_id}] [JOB-{job_id}] Restoring {len(room.cookies)} cookies")
        
        viewport_inject_code = """
import sys
_original_goto = None

def _patched_goto(self, url, **kwargs):
    result = _original_goto(self, url, **kwargs)
    try:
        self.evaluate('''
            let meta = document.querySelector('meta[name="viewport"]');
            if (!meta) {
                meta = document.createElement('meta');
                meta.name = 'viewport';
                document.head.appendChild(meta);
            }
            meta.content = 'width=1920, height=1080, initial-scale=1.0, maximum-scale=1.0, user-scalable=no';
            document.documentElement.style.width = '1920px';
            document.documentElement.style.height = '1080px';
            document.body.style.width = '1920px';
            document.body.style.height = '1080px';
        ''')
    except:
        pass
    return result

if 'browser' in dir() and browser is not None:
    try:
        from playwright.sync_api import Page
        if not hasattr(Page, '_viewport_patched'):
            _original_goto = Page.goto
            Page.goto = _patched_goto
            Page._viewport_patched = True
    except:
        pass
"""
        
        namespace = {
            'browser': browser_obj,
            'room_cookies': room.cookies,
            'public_dir': str(room.screenshots_dir),
            'room_id': room.room_id,
            'job_id': job_id,
            'time': time,
            'result': result,
            'Path': Path,
            'random': random,
            'math': math,
            'Image': Image,
            'ImageDraw': ImageDraw,
            'requests': requests,
            'print': print,
            'len': len,
            'int': int,
            'str': str,
            'dict': dict,
            'list': list,
            'VIEWPORT_WIDTH': VIEWPORT_WIDTH,
            'VIEWPORT_HEIGHT': VIEWPORT_HEIGHT,
        }
        
        exec(viewport_inject_code, namespace)
        
        print(f"[ROOM-{room.room_id}] [JOB-{job_id}] Executing code...")
        exec(code_snippet, namespace)
        
        if 'return_value' in namespace:
            result['data'] = namespace['return_value']
            
        print(f"[ROOM-{room.room_id}] [JOB-{job_id}] Execution completed successfully")
        
    except Exception as e:
        error_msg = str(e) + "\n" + traceback.format_exc()
        result['error'] = error_msg
        print(f"[ROOM-{room.room_id}] [JOB-{job_id}] ERROR: {error_msg}")
    finally:
        if browser_obj:
            try:
                browser_obj.__exit__(None, None, None)
            except:
                pass
    
    return result

def process_job(job):
    room = None
    try:
        job_manager.update_job(job.job_id, status='running', started_at=datetime.now())
        
        room = room_manager.acquire_room(job.job_id, timeout=900)
        job_manager.update_job(job.job_id, room_id=room.room_id)
        
        result = execute_in_room(job.code, room, job.job_id)
        
        screenshot_files = []
        for file in sorted(room.screenshots_dir.glob('*.png'), key=lambda x: x.stat().st_mtime):
            screenshot_files.append({
                'name': file.name,
                'publicURL': f"{SCREENSHOT_BASE_URL}/files/room_{room.room_id}/{file.name}"
            })
        
        if result.get('error'):
            job_manager.update_job(
                job.job_id,
                status='failed',
                error=result['error'],
                completed_at=datetime.now(),
                screenshots=screenshot_files
            )
        else:
            job_manager.update_job(
                job.job_id,
                status='completed',
                result=result.get('data'),
                completed_at=datetime.now(),
                screenshots=screenshot_files
            )
            
    except Exception as e:
        error_msg = str(e) + "\n" + traceback.format_exc()
        job_manager.update_job(
            job.job_id,
            status='failed',
            error=error_msg,
            completed_at=datetime.now()
        )
    finally:
        if room:
            room.cleanup_screenshots()
            room_manager.release_room(room)

@app.route('/api/s-playwright', methods=['POST'])
def execute_playwright():
    try:
        data = request.get_json()
        
        if not data or 'code' not in data or 'lang' not in data:
            return jsonify({
                'success': False,
                'error': 'Missing required fields: code and lang'
            }), 400
        
        code = data['code']
        lang = data['lang'].lower()
        
        if lang != 'python':
            return jsonify({
                'success': False,
                'error': f'Only Python is supported, got: {lang}'
            }), 400
        
        base_url = request.url_root.rstrip('/')
        job = job_manager.create_job(code, lang, base_url)
        
        thread = threading.Thread(target=process_job, args=(job,))
        thread.daemon = True
        thread.start()
        
        return jsonify({
            'success': True,
            'job_id': job.job_id,
            'status': 'queued',
            'check_url': f'/job/{job.job_id}'
        })
        
    except Exception as e:
        print(f"[API ERROR] {str(e)}\n{traceback.format_exc()}")
        return jsonify({
            'success': False,
            'error': str(e),
            'stack': traceback.format_exc()
        }), 500

@app.route('/job/<string:job_id>', methods=['GET'])
def get_job_status(job_id):
    job = job_manager.get_job(job_id.upper())
    
    if not job:
        return jsonify({
            'success': False,
            'error': 'Job not found'
        }), 404
    
    return jsonify({
        'success': True,
        'job': job.to_dict()
    })

@app.route('/files/room_<int:room_id>/<path:filename>')
def serve_file(room_id, filename):
    room_dir = PUBLIC_DIR / f'room_{room_id}'
    return send_from_directory(room_dir, filename)

@app.route('/health', methods=['GET'])
def health():
    with room_manager.lock:
        busy_rooms = sum(1 for r in room_manager.rooms.values() if r.is_busy)
        available = len(room_manager.available_rooms)
    
    with job_manager.lock:
        total_jobs = len(job_manager.jobs)
        queued_jobs = sum(1 for j in job_manager.jobs.values() if j.status == 'queued')
        running_jobs = sum(1 for j in job_manager.jobs.values() if j.status == 'running')
        
    return jsonify({
        'status': 'healthy',
        'rooms': {
            'total': MAX_ROOMS,
            'available': available,
            'busy': busy_rooms
        },
        'jobs': {
            'total': total_jobs,
            'queued': queued_jobs,
            'running': running_jobs
        },
        'viewport': {
            'width': VIEWPORT_WIDTH,
            'height': VIEWPORT_HEIGHT,
            'locked': True
        },
        'timestamp': int(time.time() * 1000)
    })

@app.route('/', methods=['GET'])
def index():
    with room_manager.lock:
        busy_rooms = sum(1 for r in room_manager.rooms.values() if r.is_busy)
        available = len(room_manager.available_rooms)
        
    return jsonify({
        'message': 'Multi-Room Camoufox Anti-Detection API with Job System',
        'endpoints': {
            'POST /api/s-playwright': 'Execute camoufox code (returns job_id)',
            'GET /job/:id': 'Check job status (3-character job ID)',
            'GET /health': 'Check API health status'
        },
        'features': [
            f'{MAX_ROOMS} Isolated Browser Rooms',
            'Job Queue System with 3-char IDs',
            'Async Job Processing',
            f'Viewport LOCKED at {VIEWPORT_WIDTH}x{VIEWPORT_HEIGHT}',
            'Camoufox Anti-Detection Browser',
            'Human-like Mouse Movement',
            'Advanced Cloudflare WAF Bypass',
            'Auto Room & Job Cleanup',
            'Concurrent Request Protection'
        ],
        'configuration': {
            'viewport': {'width': VIEWPORT_WIDTH, 'height': VIEWPORT_HEIGHT, 'locked': True},
            'port': PORT,
            'rooms': {
                'total': MAX_ROOMS,
                'available': available,
                'busy': busy_rooms
            },
            'timeout': '1200 seconds per execution',
            'room_timeout': f'{ROOM_TIMEOUT_MINUTES} minutes',
            'job_expiry': f'{JOB_EXPIRY_MINUTES} minutes',
            'screenshot_base_url': SCREENSHOT_BASE_URL
        }
    })

if __name__ == '__main__':
    print(f"🚀 Multi-Room Camoufox API with Job System")
    print(f"🌐 Port: {PORT}")
    print(f"🏠 Browser Rooms: {MAX_ROOMS}")
    print(f"📐 Viewport: LOCKED at {VIEWPORT_WIDTH}x{VIEWPORT_HEIGHT}")
    print(f"📍 Endpoints: POST /api/s-playwright, GET /job/:id")
    print(f"🎭 Features: Job queue, 3-char IDs, Async processing")
    print(f"🖱️  Human mouse movement: ENABLED")
    print(f"🌍 Proxy: http://pnode1.danbot.host:1271")
    print(f"📸 Screenshot URL: {SCREENSHOT_BASE_URL}")
    print(f"⏱️  Execution timeout: 1200 seconds")
    print(f"🔒 Room timeout: {ROOM_TIMEOUT_MINUTES} minutes")
    print(f"📦 Job expiry: {JOB_EXPIRY_MINUTES} minutes")
    
    room_manager.start_cleanup()
    
    def cleanup_jobs():
        while True:
            job_manager.cleanup_expired_jobs()
            time.sleep(300)
    
    job_cleanup_thread = threading.Thread(target=cleanup_jobs, daemon=True)
    job_cleanup_thread.start()
    
    def cleanup_on_exit():
        room_manager.stop_cleanup_thread()
        room_manager.shutdown_all()
    
    atexit.register(cleanup_on_exit)
    
    app.run(host='0.0.0.0', port=PORT, debug=False, threaded=True)