from flask import Flask, request, Response, send_from_directory, jsonify
import requests
from flask_cors import CORS
import urllib.parse
from requests.adapters import HTTPAdapter
from urllib3.util import Retry
import yt_dlp
import io

# Serve frontend static files under `/static` to avoid colliding with API routes
app = Flask(__name__, static_folder='../frontend', static_url_path='/static')
app.url_map.strict_slashes = False

# Use flask-cors for API routes
CORS(app, resources={r"/api/*": {"origins": "*"}})

# Configure a requests session with retries
session = requests.Session()
retries = Retry(total=2, backoff_factor=0.2, status_forcelist=(500, 502, 503, 504))
adapter = HTTPAdapter(max_retries=retries)
session.mount('http://', adapter)
session.mount('https://', adapter)

@app.route('/api/status', methods=['POST', 'GET'])
def proxy_status():
    if request.method == 'POST':
        # Safely get the json body
        client_data = request.get_json() or {}
        target_server = client_data.get('server')
    else:
        target_server = request.args.get('server')

    if not target_server:
        return jsonify({"error": "No server provided"}), 400
    if target_server.endswith('/'):
        target_server = target_server[:-1]
    target_server = f"{target_server.strip()}/api/status"
    # Validate target URL
    parsed = urllib.parse.urlparse(target_server)
    if parsed.scheme not in ('http', 'https') or not parsed.netloc:
        return jsonify({"error": "Invalid target URL"}), 400

    try:
        # Forward a small set of safe request headers
        safe_request_headers = {}
        for k in ('User-Agent', 'Accept', 'Accept-Encoding', 'Accept-Language', 'Range'):
            if k in request.headers:
                safe_request_headers[k] = request.headers[k]

        # Remove hop-by-hop request headers per RFC
        for hop in ('Connection', 'Keep-Alive', 'Proxy-Authenticate', 'Proxy-Authorization', 'TE', 'Trailers', 'Transfer-Encoding', 'Upgrade'):
            safe_request_headers.pop(hop, None)

        # Stream the GET via a session with small retry policy
        resp = session.get(target_server, headers=safe_request_headers, timeout=5, stream=True)

        hop_by_hop = {
            'connection', 'keep-alive', 'proxy-authenticate', 'proxy-authorization',
            'te', 'trailers', 'transfer-encoding', 'upgrade'
        }

        forwarded_headers = []
        for name, value in resp.headers.items():
            if name.lower() in hop_by_hop:
                continue
            if name.lower() == 'content-length':
                # Let Flask set Content-Length if needed for the assembled response
                continue
            forwarded_headers.append((name, value))

        def generate():
            try:
                for chunk in resp.iter_content(1024 * 16):
                    if chunk:
                        yield chunk
            finally:
                resp.close()

        return Response(generate(), status=resp.status_code, headers=forwarded_headers)
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    
@app.route('/api/videos', methods=['POST'])
def videos_proxy():
    client_data = request.get_json() or {}
    target_server = client_data.get('server')
    client_data.pop('server', None)  # Remove server from payload
    if not target_server:
        return jsonify({"error": "No server provided"}), 400
    if target_server.endswith('/'):
        target_server = target_server[:-1]
    target_server = f"{target_server.strip()}/api/videos"
    # Validate target URL
    parsed = urllib.parse.urlparse(target_server)
    if parsed.scheme not in ('http', 'https') or not parsed.netloc:
        return jsonify({"error": "Invalid target URL"}), 400

    try:
        resp = session.post(target_server, json=client_data,timeout=5)
        return Response(resp.content, status=resp.status_code, content_type=resp.headers.get('Content-Type', 'application/json'))
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/')
def index():
    return send_from_directory(app.static_folder, 'index.html')

@app.route('/api/stream', methods=['POST', 'GET'])
def stream_video():
    # Note: <video> tags perform GET. To support your POST requirement, 
    # we handle the URL via JSON post or URL params.
    video_url = ""
    if request.method == 'POST':
        video_url = request.json.get('url')
    else:
        video_url = request.args.get('url')
    
    if not video_url:
        return jsonify({"error": "No URL provided"}), 400

    def generate():
        try:
            # Configure yt-dlp options
            ydl_opts = {
                'format': 'best[ext=mp4]/best[vcodec^=avc1]/best[vcodec^=vp]/best',
                'quiet': True,
                'no_warnings': True,
                'socket_timeout': 30,
                'retries': 3,
                'fragment_retries': 3,
                'http_headers': {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                },
                'skip_unavailable_fragments': True
            }
            
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                # Extract the info
                info = ydl.extract_info(video_url, download=False)
                
                # Try to get the URL from the info dict (works for progressive downloads)
                stream_url = info.get('url')
                format_id = info.get('format_id')
                
                # If no direct URL, try to get it from formats
                if not stream_url and 'formats' in info:
                    # Find the best format that has a URL
                    for fmt in info['formats']:
                        if fmt.get('url'):
                            stream_url = fmt.get('url')
                            break
                
                if not stream_url:
                    yield b"Error: Could not extract stream URL"
                    return
                
                # Prepare headers for the stream request
                headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                }
                
                # Add any cookies or authentication headers from yt-dlp
                if 'http_headers' in info:
                    headers.update(info['http_headers'])
                
                # Stream the video from the extracted URL
                resp = session.get(stream_url, headers=headers, stream=True, timeout=30, allow_redirects=True)
                resp.raise_for_status()
                
                for chunk in resp.iter_content(chunk_size=1024 * 16):
                    if chunk:
                        yield chunk
        except Exception as e:
            yield f"Error: {str(e)}".encode()

    return Response(generate(), mimetype='video/mp4')

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