from flask import Flask, request, jsonify
import os, ai.processor
from dotenv import load_dotenv
import io
from PIL import ImageGrab

load_dotenv()

class WebServerApp:
    def __init__(self, aip):
        self.app = Flask(__name__)
        self._register_routes()
        self.aip: ai.processor.AIProcessor = aip

    def _register_routes(self):
        @self.app.route('/api/test')
        def test():
            return jsonify({"message": "Hello, World!"})

        @self.app.route('/api/request', methods=['POST'])
        def handle_request():
            # sent as form-data
            data = request.form.to_dict()
            if not data:
                return jsonify({"error": "No data provided"}), 400
            
            # Process the data as needed
            prompt = data.get('prompt', '')



            if not prompt:
                return jsonify({"error": "No prompt provided"}), 400
            img_data = None
            if 'img' in request.files:
                img_file = request.files['img']
                if img_file:
                    img_data = img_file.read()
                else:
                    img_data = None
            else:
                if 'host_screenshot' in data:
                    # take a screenshot right here
                    # capture the full screen
                    screenshot_img = ImageGrab.grab()
                    buf = io.BytesIO()
                    screenshot_img.save(buf, format='PNG')
                    img_data = buf.getvalue() 
            
            import base64
            # Convert image data to base64 if provided
            if img_data and isinstance(img_data, bytes):
                img_data = base64.b64encode(img_data).decode('utf-8')
            elif img_data and isinstance(img_data, str):
                img_data = img_data.encode('utf-8')                
            
            response = self.aip.process(prompt, img_data)
            return jsonify({"response": response}), 200

        @self.app.route('/api/health')
        def health():
            return jsonify({"status": "healthy"})

    def run(self, *args, **kwargs):
        self.app.run(*args, **kwargs)

# Example usage:
# if __name__ == "__main__":
#     server = WebServerApp()
#     server.run()