from flask import Blueprint, render_template, request, jsonify, flash, redirect, url_for, current_app, Response
from flask_login import login_required, current_user
from websockets.sync.client import connect as ws_connect
import websockets.exceptions
import threading
import requests
from urllib.parse import urlparse
from app.models import SystemConfig, db
from app.utils.octoprint_client import OctoPrintClient
from app.models import PrintFile
import os

printer_bp = Blueprint('printer', __name__, url_prefix='/printer')

def get_octo_client():
    url = SystemConfig.query.filter_by(key='octoprint_url').first()
    apikey = SystemConfig.query.filter_by(key='octoprint_apikey').first()
    if url and url.value and apikey and apikey.value:
        return OctoPrintClient(url.value, apikey.value)
    return None

@printer_bp.route('/status')
@login_required
def status():
    client = get_octo_client()
    status_data = None
    job_data = None
    error = None
    if client:
        try:
            status_data = client.get_printer_status()
            job_data = client.get_job_info()
        except Exception as e:
            error = str(e)
    else:
        error = "OctoPrint is not configured."
        
    return render_template('printer/status.html', status=status_data, job=job_data, error=error)

def get_gcode_dir():
    conf = SystemConfig.query.filter_by(key='gcode_upload_folder').first()
    if conf and conf.value and os.path.exists(conf.value):
        return conf.value
    return current_app.config['UPLOAD_FOLDER']

@printer_bp.route('/prepare')
@login_required
def prepare():
    from app.models import PrintFile
    import os
    
    # Query only the sliced GCode files belonging to the current user
    user_files = PrintFile.query.filter_by(user_id=current_user.id, status='sliced').order_by(PrintFile.created_at.desc()).all()
    
    files = []
    gcode_dir = get_gcode_dir()
    
    client = get_octo_client()
    octo_files_dict = {}
    if client:
        try:
            octo_resp = client.get_files()
            for item in octo_resp.get('files', []):
                octo_files_dict[item.get('name')] = item
        except Exception as e:
            pass
            
    for f in user_files:
        gcode_filename = f.filename.rsplit('.', 1)[0] + '.gcode'
        gcode_path = os.path.join(gcode_dir, gcode_filename)
        
        size = 0
        if os.path.exists(gcode_path):
            size = os.path.getsize(gcode_path)
            
        # Upload to OctoPrint if not found but exists locally
        if client and gcode_filename not in octo_files_dict and size > 0:
            try:
                resp = client.upload_file('local', gcode_path, gcode_filename)
                uploaded_loc = resp.get('files', {}).get('local', {})
                if gcode_filename in uploaded_loc:
                    octo_files_dict[gcode_filename] = uploaded_loc[gcode_filename]
            except Exception as e:
                pass
                
        octo_info = octo_files_dict.get(gcode_filename, {})
        analysis = octo_info.get('gcodeAnalysis', None)
        
        files.append({
            'id': f.id,
            'name': f.original_filename.rsplit('.', 1)[0] + '.gcode',
            'type': 'machinecode',
            'size': size,
            'origin': 'local',
            'path': gcode_filename,
            'gcodeAnalysis': analysis
        })
        
    error = None
    if not get_octo_client():
        error = "OctoPrint is not configured."
        
    return render_template('printer/prepare.html', files=files, error=error)

@printer_bp.route('/api/print_file', methods=['POST'])
@login_required
def api_print_file():
    path = request.json.get('path')
    location = request.json.get('origin', 'local')
    client = get_octo_client()
    if client and path:
        try:
            client.select_file(location, path, print_after_select=True)
            return jsonify({"success": True})
        except Exception as e:
            return jsonify({"success": False, "error": str(e)})
    return jsonify({"success": False, "error": "Not configured or missing path"})

@printer_bp.route('/control')
@login_required
def control():
    client = get_octo_client()
    webcam_url = None
    error = None
    if client:
        try:
            webcam_url = client.get_webcam_stream_url()
        except Exception as e:
            error = str(e)
    else:
        error = "OctoPrint is not configured."
    return render_template('printer/control.html', webcam_url=webcam_url, error=error)

@printer_bp.route('/api/command', methods=['POST'])
@login_required
def api_command():
    cmd = request.json.get('command')
    client = get_octo_client()
    if client and cmd:
        try:
            if cmd == 'home':
                client.home_axes()
            elif cmd == 'pause':
                client.pause_print()
            elif cmd == 'cancel':
                client.cancel_print()
            return jsonify({"success": True})
        except Exception as e:
            return jsonify({"success": False, "error": str(e)})
    return jsonify({"success": False, "error": "Invalid client or command"})

@printer_bp.route('/api/upload_gcode', methods=['POST'])
@login_required
def upload_gcode():
    from app.models import PrintFile
    import os
    import uuid
    from werkzeug.utils import secure_filename
    
    if 'file' not in request.files:
        return jsonify({"success": False, "error": "No file part"}), 400
        
    file = request.files['file']
    if file.filename == '':
        return jsonify({"success": False, "error": "No selected file"}), 400
        
    if not file.filename.lower().endswith(('.gcode', '.gco', '.g')):
        return jsonify({"success": False, "error": "Only standard .gcode files are supported."}), 400
        
    sec_name = secure_filename(file.filename)
    random_prefix = uuid.uuid4().hex[:8]
    
    # We save pseudo-STL filename so that our conventional parser works (replaces .ext with .gcode)
    # i.e., "some_file.gcode" -> pseudo .stl tracker
    pseudo_stl_filename = f"{random_prefix}_{sec_name.rsplit('.', 1)[0]}.stl"
    gcode_filename = f"{random_prefix}_{sec_name.rsplit('.', 1)[0]}.gcode"
    
    gcode_path = os.path.join(get_gcode_dir(), gcode_filename)
    file.save(gcode_path)
    
    print_file = PrintFile(
        user_id=current_user.id,
        original_filename=file.filename, # keep original GCode name
        filename=pseudo_stl_filename,
        file_type='stl',
        status='sliced'
    )
    db.session.add(print_file)
    db.session.commit()
    
    client = get_octo_client()
    if client:
        try:
            client.upload_file('local', gcode_path, gcode_filename)
        except Exception:
            pass
            
    return jsonify({"success": True})

@printer_bp.route('/octo_config', methods=['GET', 'POST'])
@login_required
def octo_config():
    if not current_user.is_admin:
        flash("Admin access required", "danger")
        return redirect(url_for('printer.status'))
        
    if request.method == 'POST':
        url = request.form.get('octoprint_url', '')
        apikey = request.form.get('octoprint_apikey', '')
        
        conf_url = SystemConfig.query.filter_by(key='octoprint_url').first()
        if not conf_url:
            conf_url = SystemConfig(key='octoprint_url')
            db.session.add(conf_url)
        conf_url.value = url.rstrip('/')
        
        conf_key = SystemConfig.query.filter_by(key='octoprint_apikey').first()
        if not conf_key:
            conf_key = SystemConfig(key='octoprint_apikey')
            db.session.add(conf_key)
        conf_key.value = apikey
        
        db.session.commit()
        if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
            return jsonify({'success': True}), 200
        flash("OctoPrint settings updated", "success")
        return redirect(url_for('printer.octo_config'))
        
    configs = {c.key: c.value for c in SystemConfig.query.all()}
    return render_template('printer/octo_config.html', configs=configs)

@printer_bp.route('/octo_embed')
@login_required
def octo_embed():
    if not current_user.is_admin:
        flash("Admin access required", "danger")
        return redirect(url_for('printer.status'))
        
    url = SystemConfig.query.filter_by(key='octoprint_url').first()
    embed_url = url_for('printer.octo_proxy') if url and url.value else None
    return render_template('printer/octo_embed.html', embed_url=embed_url)

@printer_bp.route('/proxy', defaults={'path': ''}, websocket=True, strict_slashes=False)
@printer_bp.route('/proxy/<path:path>', websocket=True)
@printer_bp.route('/proxy', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'], strict_slashes=False)
@printer_bp.route('/proxy/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
@login_required
def octo_proxy(path):
    if not current_user.is_admin:
        return "Unauthorized", 403

    url_config = SystemConfig.query.filter_by(key='octoprint_url').first()
    if not url_config or not url_config.value:
        return "OctoPrint URL not configured", 404

    base_url = url_config.value.rstrip('/')
    
    # print("----- REQUEST HEADERS -----")
    # for k, v in request.headers:
    #     print(f"{k}: {v}")
    # print("----- END REQUEST HEADERS -----")

    # --- WebSocket Proxy Logic ---
    if request.headers.get('Upgrade', '').lower() == 'websocket':
        from flask_sock import Server, ConnectionClosed
        
        # Check if environment supports WebSockets
        try:
            ws = Server(request.environ)
        except Exception as e:
            env_keys = sorted(list(request.environ.keys()))
            print(f"FAILED. ENV KEYS: {env_keys}")
            return f"WebSocket Upgrade Failed: {str(e)}", 400
            
        def handle_ws():
            if base_url.startswith('https://'):
                ws_base = base_url.replace('https://', 'wss://', 1)
            else:
                ws_base = base_url.replace('http://', 'ws://', 1)
                
            target_url = f"{ws_base}/{path}"
            if request.query_string:
                target_url = f"{target_url}?{request.query_string.decode('utf-8')}"
                print(f"WS Proxy Query String: {request.query_string.decode('utf-8')}")
            # Copy most headers, especially Origin to pass CORS
            ws_headers = {}
            for k, v in request.headers:
                if k.lower() not in ['host', 'connection', 'upgrade', 'sec-websocket-key', 'sec-websocket-version', 'sec-websocket-extensions', 'content-length']:
                    ws_headers[k] = v
            
            # Match Tornado's expectations for Origin to avoid 400 Bad Request
            parsed_base = urlparse(base_url)
            ws_headers['Host'] = parsed_base.netloc
            if 'Origin' in request.headers:
                ws_headers['Origin'] = base_url
            if 'Referer' in request.headers:
                ws_headers['Referer'] = f"{base_url}/{path}"
                
            ws_headers['X-Real-IP'] = request.remote_addr
            forwarded_for = request.headers.get('X-Forwarded-For', '')
            ws_headers['X-Forwarded-For'] = f"{forwarded_for}, {request.remote_addr}" if forwarded_for else request.remote_addr
            ws_headers['X-Forwarded-Proto'] = request.scheme
            ws_headers['X-Forwarded-Host'] = request.host
            ws_headers['X-Forwarded-Port'] = str(request.environ.get('SERVER_PORT', '80'))
            ws_headers['X-Script-Name'] = '/printer/proxy'
            
            print(f"WS Proxy Connecting to: {target_url}")
            try:
                remote_ws = ws_connect(target_url, additional_headers=ws_headers)
                print("WS Proxy connected to remote.")
            except Exception as e:
                import traceback
                traceback.print_exc()
                print(f"Remote WS Connection Error: {e}")
                ws.close(1011, str(e))
                return

            def recv_loop():
                print("WS recv_loop started")
                try:
                    for message in remote_ws:
                        ws.send(message)
                except Exception as e:
                    print("WS recv error:", e)
                finally:
                    try: remote_ws.close()
                    except: pass
                    try: ws.close()
                    except: pass
                print("WS recv_loop ended")

            t = threading.Thread(target=recv_loop)
            t.daemon = True
            t.start()

            print("WS Entering client receive loop")
            try:
                while True:
                    data = ws.receive()
                    if data is None:
                        break
                    remote_ws.send(data)
            except Exception as e:
                print("WS send error:", e)
            finally:
                try: remote_ws.close()
                except: pass
                try: ws.close()
                except: pass
            print("WS client loop ended")
        
        try:
            handle_ws()
        except ConnectionClosed:
            print("WS Connection Closed")
        except Exception as e:
            print("WS Error in handle_ws:", e)
        finally:
            try: ws.close() 
            except: pass

        class WebSocketResponse(Response):
            def __call__(self, *args, **kwargs):
                print("WS Response __call__")
                if getattr(ws, 'mode', 'werkzeug') == 'werkzeug':
                    return super().__call__(*args, **kwargs)
                return []
                
        return WebSocketResponse()

    # --- Standard HTTP Proxy Logic ---
    # from urllib.parse import urlparse
    target_url = f"{base_url}/{path}"
    
    if request.query_string:
        target_url = f"{target_url}?{request.query_string.decode('utf-8')}"

    # Build headers for reverse proxy based on nginx config reference
    parsed_base = urlparse(base_url)
    headers = {k: v for k, v in request.headers if k.lower() not in ['host', 'content-length', 'origin', 'referer']}
    
    # print(f"Proxying to: {target_url}")
    # Spoof Host, Origin, and Referer to match the backend URL completely
    # This prevents Tornado's strict Origin vs Host CSRF/CORS validation from failing
    # headers['Host'] = parsed_base.netloc
    headers['Host'] = request.host
    
    if 'Origin' in request.headers:
        headers['Origin'] = base_url
    if 'Referer' in request.headers:
        headers['Referer'] = f"{base_url}/{path}"
    
    headers['X-Real-IP'] = request.remote_addr
    headers['X-Real-Port'] = str(request.environ.get('REMOTE_PORT', ''))
    
    forwarded_for = request.headers.get('X-Forwarded-For', '')
    headers['X-Forwarded-For'] = f"{forwarded_for}, {request.remote_addr}" if forwarded_for else request.remote_addr
    
    headers['X-Forwarded-Proto'] = request.scheme
    headers['X-Script-Name'] = "/printer/proxy"
    headers['X-Forwarded-Host'] = request.host
    headers['X-Forwarded-Port'] = str(request.environ.get('SERVER_PORT', '80'))
    headers['REMOTE-HOST'] = request.remote_addr
    
    if request.headers.get('Upgrade'):
        headers['Upgrade'] = request.headers.get('Upgrade')
    if request.headers.get('Connection'):
        headers['Connection'] = request.headers.get('Connection')
        
    # OctoPrint requires an X-CSRF-Token header matching the csrf_token_* cookie for POST/PUT/DELETE
    if request.method not in ['GET', 'HEAD', 'OPTIONS'] and 'X-CSRF-Token' not in request.headers:
        for cookie_name, cookie_value in request.cookies.items():
            if cookie_name.startswith('csrf_token_'):
                headers['X-CSRF-Token'] = cookie_value
                break

    # if path == 'api/login':
    #     print("----- SEND HEADERS -----")
    #     for k, v in headers.items():
    #         if k in request.headers.keys():
    #             print(f"{k} :(from request): {request.headers[k]} (to send): {v}")
    #         else:
    #             print(f"{k} :(from request): None (to send): {v}")
        
    #     for k,v in request.headers.items():
    #         if k not in headers.keys():
    #             print(f"{k} :(from request): {request.headers[k]} (to send): None")
        
    #     print("----- END SEND HEADERS -----")

    try:
        # proxy_connect_timeout 60s, proxy_read_timeout 600s
        resp = requests.request(
            method=request.method,
            url=target_url,
            headers=headers,
            data=request.get_data(),
            cookies=request.cookies,
            allow_redirects=False,
            stream=True,
            timeout=(60, 600)
        )
    except requests.exceptions.RequestException as e:
        return f"Proxy connection error: {str(e)}", 502

    # Strip headers that might break the iframe or framing
    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection', 'x-frame-options', 'content-security-policy']
    
    # We must use raw headers to prevent requests from joining multiple Set-Cookie headers with a comma
    # Joining Set-Cookie with a comma breaks standard cookie parsing in the browser due to commas in dates
    response_headers = [(name, value) for name, value in resp.raw.headers.items() 
                        if name.lower() not in excluded_headers and name.lower() != 'set-cookie']
    
    for cookie in resp.raw.headers.get_all('set-cookie', []):
        # ensure we preserve the proxy path override
        response_headers.append(('Set-Cookie', cookie))

    def generate():
        for chunk in resp.iter_content(chunk_size=8192):
            if chunk:
                yield chunk

    return Response(generate(), resp.status_code, response_headers)