#!/usr/bin/env python3
"""
Jellyfin Xbox Codec Filter Proxy
Filters AV1 and Opus from DeviceProfile responses for Xbox clients
"""
import json
import logging
from flask import Flask, request, Response
import requests

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

JELLYFIN_URL = "http://jellyfin:8096"

def filter_codecs(profile):
    """Remove AV1 from video codecs and Opus from audio codecs"""
    if not profile:
        return profile

    # Filter DirectPlayProfiles
    if profile.get("DirectPlayProfiles"):
        for direct_play in profile["DirectPlayProfiles"]:
            if "VideoCodec" in direct_play and direct_play["VideoCodec"]:
                codecs = [c.strip() for c in direct_play["VideoCodec"].split(",")]
                codecs = [c for c in codecs if c.lower() != "av1"]
                direct_play["VideoCodec"] = ",".join(codecs)

            if "AudioCodec" in direct_play and direct_play["AudioCodec"]:
                codecs = [c.strip() for c in direct_play["AudioCodec"].split(",")]
                codecs = [c for c in codecs if c.lower() != "opus"]
                direct_play["AudioCodec"] = ",".join(codecs)

    # Filter TranscodingProfiles
    if profile.get("TranscodingProfiles"):
        for transcoding in profile["TranscodingProfiles"]:
            if "VideoCodec" in transcoding and transcoding["VideoCodec"]:
                codecs = [c.strip() for c in transcoding["VideoCodec"].split(",")]
                codecs = [c for c in codecs if c.lower() != "av1"]
                transcoding["VideoCodec"] = ",".join(codecs)

            if "AudioCodec" in transcoding and transcoding["AudioCodec"]:
                codecs = [c.strip() for c in transcoding["AudioCodec"].split(",")]
                codecs = [c for c in codecs if c.lower() != "opus"]
                transcoding["AudioCodec"] = ",".join(codecs)

    return profile

@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
def proxy(path):
    """Proxy all requests to Jellyfin, filter Xbox responses"""
    user_agent = request.headers.get('User-Agent', '')
    is_xbox = 'xbox' in user_agent.lower()
    is_playback_info = 'PlaybackInfo' in path and request.method == 'POST'

    # Build full URL with query parameters
    url = f"{JELLYFIN_URL}/{path}"
    if request.query_string:
        url += f"?{request.query_string.decode('utf-8')}"

    # Copy headers, exclude problematic ones
    headers = {}
    excluded = ['host', 'connection', 'content-length', 'content-encoding', 'transfer-encoding']
    for key, value in request.headers:
        if key.lower() not in excluded:
            headers[key] = value

    # Add X-Forwarded-* headers for proper client tracking
    headers['X-Forwarded-For'] = request.remote_addr
    headers['X-Forwarded-Proto'] = request.scheme
    headers['X-Forwarded-Host'] = request.host
    headers['X-Real-IP'] = request.remote_addr

    # Get request body
    request_data = request.get_data()

    # Filter REQUEST body for Xbox PlaybackInfo (DeviceProfile is sent TO server)
    if is_xbox and is_playback_info:
        logger.info(f"Xbox PlaybackInfo request detected: {path}, User-Agent: {user_agent}")
        if request_data:
            try:
                data = json.loads(request_data)
                if "DeviceProfile" in data:
                    logger.info(f"Filtering Xbox DeviceProfile in REQUEST")
                    original_profile = data["DeviceProfile"]
                    filtered_profile = filter_codecs(data["DeviceProfile"])
                    data["DeviceProfile"] = filtered_profile
                    request_data = json.dumps(data).encode('utf-8')
                    headers['Content-Length'] = str(len(request_data))
                    logger.info("DeviceProfile filtered successfully in REQUEST")
                else:
                    logger.warning("No DeviceProfile found in REQUEST body")
            except Exception as e:
                logger.error(f"Error filtering request: {e}", exc_info=True)
        else:
            logger.warning("No request body found for PlaybackInfo")

    resp = requests.request(
        method=request.method,
        url=url,
        headers=headers,
        data=request_data,
        cookies=request.cookies,
        allow_redirects=False,
        stream=True  # Support large file streaming
    )

    # Stream all responses (no need to filter response, DeviceProfile only in request)
    response = Response(resp.iter_content(chunk_size=8192), status=resp.status_code)

    # Copy headers from Jellyfin response
    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
    for key, value in resp.headers.items():
        if key.lower() not in excluded_headers:
            response.headers[key] = value

    return response

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8096)
