#!/usr/bin/env python3
import http.server
import http.cookiejar
import urllib.request
import urllib.error
import urllib.parse
import json
import re
import base64
import time
import sys
import os
import argparse
import threading
import traceback
import gc
from collections import OrderedDict
from socketserver import ThreadingMixIn
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from concurrent.futures import ThreadPoolExecutor

CONFIG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'cache', 'orig_urls.json')
PROXY_PORT = 1338
PROXY_BIND = '0.0.0.0'
UA = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36'

MAX_COOKIES_PER_JAR = 50
MAX_CACHED_OPENERS = 20
MANIFEST_CACHE_SIZE = 100
MAX_THREAD_POOL = 32
GC_INTERVAL_SEC = 60
SEGMENT_STREAM_BUF = 65536

# Global default auth (user, pass)
_global_auth = None

# Cached channel URLs (loaded once)
_channel_urls = {}
_channel_urls_lock = threading.Lock()

def _load_channel_urls():
    global _channel_urls
    with _channel_urls_lock:
        if not _channel_urls:
            with open(CONFIG_FILE) as f:
                raw = json.load(f)
            # Support both {"name": "url"} and {"name": {"url":"...","auth":[...]}}
            for k, v in raw.items():
                _channel_urls[k] = v
            log(f"Loaded {len(_channel_urls)} channel URLs")
    return _channel_urls

def _get_channel_entry(name):
    """Return (url, auth) tuple for a channel. auth may be None."""
    urls = _load_channel_urls()
    entry = urls.get(name)
    if entry is None:
        return None, None
    if isinstance(entry, dict):
        url = entry.get('url', '')
        auth = entry.get('auth')  # list [user, pass] or None
        return url, auth
    return entry, None  # plain string URL

# Bounded CookieJar
class BoundedCookieJar(http.cookiejar.CookieJar):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def set_cookie(self, cookie):
        super().set_cookie(cookie)
        if len(self) > MAX_COOKIES_PER_JAR:
            to_remove = len(self) - MAX_COOKIES_PER_JAR
            cookies = list(self)
            for c in cookies[:to_remove]:
                self.clear(c.domain, c.path, c.name)

# Session/Cookie management
_sessions = {}
_sessions_lock = threading.Lock()
_session_last_access = {}

def _get_jar(domain):
    with _sessions_lock:
        if domain not in _sessions:
            _sessions[domain] = BoundedCookieJar()
        _session_last_access[domain] = time.time()
        return _sessions[domain]

def _evict_stale_sessions(max_age=3600):
    with _sessions_lock:
        now = time.time()
        stale = [d for d, t in _session_last_access.items() if now - t > max_age]
        for d in stale:
            del _sessions[d]
            del _session_last_access[d]
        if stale:
            log(f"Evicted {len(stale)} stale sessions, {len(_sessions)} remaining")

# Cached openers (LRU)
_opener_cache = OrderedDict()
_opener_cache_lock = threading.Lock()

def _get_opener(cookie_jar=None):
    jar_id = id(cookie_jar) if cookie_jar else 0
    with _opener_cache_lock:
        if jar_id in _opener_cache:
            _opener_cache.move_to_end(jar_id)
            return _opener_cache[jar_id]
    handlers = []
    if cookie_jar:
        handlers.append(urllib.request.HTTPCookieProcessor(cookie_jar))
    opener = urllib.request.build_opener(*handlers)
    with _opener_cache_lock:
        _opener_cache[jar_id] = opener
        _opener_cache.move_to_end(jar_id)
        while len(_opener_cache) > MAX_CACHED_OPENERS:
            _opener_cache.popitem(last=False)
    return opener

# URL fetching — returns (body, status, content_type, final_url)
def fetch_url(url, timeout=15, cookie_jar=None, referer=None, auth=None):
    """Fetch a URL, following redirects. Returns (body, status, content_type, final_url).
    
    final_url is the URL after any HTTP redirects, which must be used as the
    base for resolving relative URLs in HLS manifests.
    """
    opener = _get_opener(cookie_jar)
    headers = {
        'User-Agent': UA,
        'Accept-Encoding': 'identity',  # prevent transparent gzip to avoid double-decode issues
    }
    if referer:
        headers['Referer'] = referer
    # Auth: per-request > global default
    effective_auth = auth or _global_auth
    if effective_auth:
        if isinstance(effective_auth, (list, tuple)) and len(effective_auth) == 2:
            cred = base64.b64encode(f"{effective_auth[0]}:{effective_auth[1]}".encode()).decode()
            headers['Authorization'] = f"Basic {cred}"
        elif isinstance(effective_auth, str) and ':' in effective_auth:
            cred = base64.b64encode(effective_auth.encode()).decode()
            headers['Authorization'] = f"Basic {cred}"
    req = urllib.request.Request(url, headers=headers)
    try:
        with opener.open(req, timeout=timeout) as resp:
            # Capture the final URL after any redirects
            final_url = getattr(resp, 'url', None) or getattr(resp, 'geturl', lambda: url)() or url
            return resp.read(), resp.status, resp.headers.get('Content-Type', 'application/octet-stream'), final_url
    except urllib.error.HTTPError as e:
        body = e.read() if e.fp else b''
        ct = e.headers.get('Content-Type', 'application/octet-stream') if e.headers else 'text/html'
        final_url = getattr(e, 'url', None) or url
        return body, e.code, ct, final_url
    except (urllib.error.URLError, OSError) as e:
        return b'', 0, 'text/plain', url

# Manifest URL rewriting
_compiled_uri_re = re.compile(r'URI="([^"]+)"')

def rewrite_manifest(body, base_url):
    """Rewrite all relative URLs in an HLS manifest to absolute proxy URLs."""
    lines = body.split('\n')
    rewritten = []
    for line in lines:
        stripped = line.strip()
        if not stripped:
            rewritten.append(line)
            continue
        # Handle #EXT-X-* tags that contain URI="..."
        if stripped.startswith('#EXT-X-') and 'URI="' in stripped:
            rewritten.append(_compiled_uri_re.sub(
                lambda m: f'URI="{_abs(m.group(1), base_url)}"', stripped))
        elif stripped.startswith('#'):
            # Other comment/tag lines — keep as-is
            rewritten.append(line)
        elif not stripped.startswith('#'):
            # Non-comment, non-empty line — treat as a resource URL
            rewritten.append(_abs(stripped, base_url))
        else:
            rewritten.append(line)
    return '\n'.join(rewritten)

def _abs(url, base):
    if url.startswith(f'http://127.0.0.1:{PROXY_PORT}') or url.startswith('data:'):
        return url
    resolved = url if url.startswith('http') else urllib.parse.urljoin(base, url)
    return _to_proxy(resolved, base)

def _to_proxy(full_url, referer):
    encoded = base64.urlsafe_b64encode(full_url.encode()).decode().rstrip('=')
    route = '/m/' if urlparse(full_url).path.lower().endswith('.m3u8') else '/s/'
    proxy_url = f"http://127.0.0.1:{PROXY_PORT}{route}{encoded}"
    if referer:
        ref_enc = base64.urlsafe_b64encode(referer.encode()).decode().rstrip('=')
        proxy_url += f"?r={ref_enc}"
    return proxy_url

def _from_proxy(encoded):
    pad = 4 - len(encoded) % 4
    if pad != 4:
        encoded += '=' * pad
    return base64.urlsafe_b64decode(encoded.encode()).decode()

def _get_referer(query_string):
    if not query_string:
        return None
    qs = parse_qs(query_string)
    if 'r' in qs:
        return _from_proxy(qs['r'][0])
    return None

# Stats
_stats = {
    'requests': 0,
    'manifests': 0,
    'segments': 0,
    'bytes_sent': 0,
    'errors': 0,
}
_stats_lock = threading.Lock()

# Request Handler
class RobustHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        with _stats_lock:
            _stats['requests'] += 1
        parsed = urlparse(self.path)
        path = parsed.path
        query = parsed.query
        try:
            if path.startswith('/channel/'):
                self._handle_channel(path, query)
            elif path.startswith('/m/'):
                self._handle_manifest(path, query)
            elif path.startswith('/s/'):
                self._handle_segment(path, query)
            else:
                self._err(404, 'Not found')
        except Exception as e:
            with _stats_lock:
                _stats['errors'] += 1
            log(f"Handler error: {traceback.format_exc()}")
            try:
                self._err(500, str(e))
            except:
                pass

    def _handle_channel(self, path, query):
        parts = path.split('/')
        if len(parts) >= 4 and parts[3] == 'master.m3u8':
            name = urllib.parse.unquote(parts[2])
            url, auth = _get_channel_entry(name)
            if url:
                q = query
                if q:
                    url += ('&' if '?' in url else '?') + q
                jar = _get_jar(urlparse(url).netloc)
                body, st, ct, final_url = fetch_url(url, cookie_jar=jar, auth=auth)
                if st == 200:
                    with _stats_lock:
                        _stats['manifests'] += 1
                    # Use final_url (after redirects) as the base for resolving relative URLs
                    text = body.decode('utf-8', 'replace')
                    # Strip BOM if present
                    if text and text[0] == '\ufeff':
                        text = text[1:]
                    rw = rewrite_manifest(text, final_url)
                    log(f"[channel/{name}] base={url} final={final_url}")
                    self._ok('application/vnd.apple.mpegurl', rw.encode())
                else:
                    log(f"[channel/{name}] upstream returned {st} for {url}")
                    self._err(502, f"Upstream {st}")
            else:
                self._err(404, "Channel not found")
        else:
            self._err(400, "Bad channel path")

    def _handle_manifest(self, path, query):
        target = _from_proxy(path[3:])
        referer = _get_referer(query)
        jar = _get_jar(urlparse(target).netloc)
        body, st, ct, final_url = fetch_url(target, cookie_jar=jar, referer=referer)
        if st == 200:
            with _stats_lock:
                _stats['manifests'] += 1
            # Use final_url (after redirects) as the base for resolving relative URLs
            text = body.decode('utf-8', 'replace')
            if text and text[0] == '\ufeff':
                text = text[1:]
            rw = rewrite_manifest(text, final_url)
            self._ok('application/vnd.apple.mpegurl', rw.encode())
        else:
            log(f"[manifest] upstream returned {st} for {target}")
            self._err(502, f"Upstream {st}")

    def _handle_segment(self, path, query):
        target = _from_proxy(path[3:])
        if target.startswith('data:'):
            cp = target.find(',')
            if cp >= 0:
                meta, payload = target[5:cp], target[cp + 1:]
                b = base64.b64decode(payload) if 'base64' in meta else urllib.parse.unquote_to_bytes(payload)
                self._ok('application/octet-stream', b)
                return
        referer = _get_referer(query)
        jar = _get_jar(urlparse(target).netloc)
        headers = {
            'User-Agent': UA,
            'Accept-Encoding': 'identity',
        }
        if referer:
            headers['Referer'] = referer
        if _global_auth:
            cred = base64.b64encode(_global_auth.encode()).decode()
            headers['Authorization'] = f"Basic {cred}"
        opener = _get_opener(jar)
        req = urllib.request.Request(target, headers=headers)
        try:
            with opener.open(req, timeout=30) as resp:
                ct = resp.headers.get('Content-Type', 'application/octet-stream')
                content_length = resp.headers.get('Content-Length')
                self.send_response(200)
                self.send_header('Content-Type', ct)
                if content_length:
                    self.send_header('Content-Length', content_length)
                self.send_header('Access-Control-Allow-Origin', '*')
                self.end_headers()
                total_sent = 0
                while True:
                    chunk = resp.read(SEGMENT_STREAM_BUF)
                    if not chunk:
                        break
                    try:
                        self.wfile.write(chunk)
                        total_sent += len(chunk)
                    except (BrokenPipeError, ConnectionResetError, OSError):
                        break
                with _stats_lock:
                    _stats['segments'] += 1
                    _stats['bytes_sent'] += total_sent
        except urllib.error.HTTPError as e:
            self._err(502, f'Upstream {e.code}')
        except (urllib.error.URLError, OSError) as e:
            self._err(502, f'Fetch error: {type(e).__name__}')
        except Exception as e:
            try:
                self._err(502, str(e))
            except:
                pass

    def _ok(self, ct, body):
        try:
            self.send_response(200)
            self.send_header('Content-Type', ct)
            self.send_header('Content-Length', str(len(body)))
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()
            self.wfile.write(body)
            with _stats_lock:
                _stats['bytes_sent'] += len(body)
        except (BrokenPipeError, ConnectionResetError, OSError):
            pass

    def _err(self, code, msg):
        try:
            self.send_response(code)
            self.send_header('Content-Type', 'text/plain')
            body = msg.encode()
            self.send_header('Content-Length', str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        except (BrokenPipeError, ConnectionResetError, OSError):
            pass

    def log_message(self, *a):
        pass

    def handle(self):
        try:
            BaseHTTPRequestHandler.handle(self)
        except (BrokenPipeError, ConnectionResetError, OSError, ValueError):
            pass
        except Exception:
            pass

# Server with bounded thread pool
class BoundedThreadedServer(ThreadingMixIn, HTTPServer):
    daemon_threads = True
    allow_reuse_address = True
    request_queue_size = 128

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._thread_pool = ThreadPoolExecutor(max_workers=MAX_THREAD_POOL)

    def process_request(self, request, client_address):
        self._thread_pool.submit(self._process_request_thread, request, client_address)

    def _process_request_thread(self, request, client_address):
        try:
            self.finish_request(request, client_address)
        except Exception:
            self.handle_error(request, client_address)
        finally:
            self.shutdown_request(request)

    def handle_error(self, request, client_address):
        pass

# Periodic maintenance
def _maintenance_loop():
    while True:
        time.sleep(GC_INTERVAL_SEC)
        collected = gc.collect()
        _evict_stale_sessions()
        pid = os.getpid()
        rss = _get_rss_kb(pid)
        with _stats_lock:
            s = dict(_stats)
        log(f"STATS: rss={rss}KB reqs={s['requests']} manifests={s['manifests']} "
            f"segments={s['segments']} bytes={s['bytes_sent']//1024//1024}MB "
            f"errors={s['errors']} sessions={len(_sessions)} openers={len(_opener_cache)} "
            f"gc_collected={collected}")

def _get_rss_kb(pid):
    try:
        with open(f'/proc/{pid}/status') as f:
            for line in f:
                if line.startswith('VmRSS:'):
                    return int(line.split()[1])
    except:
        return 0

def log(msg):
    sys.stderr.write(f"[{time.strftime('%H:%M:%S')}] {msg}\n")
    sys.stderr.flush()

# Main
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='HLS Proxy for o11pro (redirect-aware + auth)')
    parser.add_argument('--config', default=CONFIG_FILE,
                        help=f'Path to channel URLs JSON (default: {CONFIG_FILE})')
    parser.add_argument('--port', type=int, default=PROXY_PORT,
                        help=f'Proxy listen port (default: {PROXY_PORT})')
    parser.add_argument('--bind', default=PROXY_BIND,
                        help=f'Proxy bind address (default: {PROXY_BIND})')
    parser.add_argument('--auth', default=None,
                        help='Global Basic Auth credentials as user:pass '
                             '(e.g. 1acd2d7afd8cb34099cb832862e3c08d:b85491a27c2852323ac704a07cf7b779)')
    args = parser.parse_args()

    CONFIG_FILE = args.config
    PROXY_PORT = args.port
    PROXY_BIND = args.bind

    if args.auth:
        _global_auth = args.auth
        log(f"Global auth enabled")

    _load_channel_urls()
    assert _compiled_uri_re is not None

    maint = threading.Thread(target=_maintenance_loop, daemon=True)
    maint.start()

    print(f"HLS Proxy on {PROXY_BIND}:{PROXY_PORT} (redirect-aware + auth)", flush=True)
    print(f"  Config:       {CONFIG_FILE}", flush=True)
    print(f"  Channels:     {len(_channel_urls)}", flush=True)
    print(f"  Thread pool:  {MAX_THREAD_POOL} workers", flush=True)
    print(f"  Cookie limit: {MAX_COOKIES_PER_JAR}/jar", flush=True)
    print(f"  Streaming buffer: {SEGMENT_STREAM_BUF} bytes", flush=True)
    print(f"  GC interval:  {GC_INTERVAL_SEC}s", flush=True)
    if _global_auth:
        print(f"  Global auth:   enabled", flush=True)
    print(f"", flush=True)
    print(f"  Routes:", flush=True)
    print(f"    /channel/{{name}}/master.m3u8  - Channel master manifest", flush=True)
    print(f"    /m/{{b64_url}}                 - Manifest fetch+rewrite", flush=True)
    print(f"    /s/{{b64_url}}                 - Segment stream", flush=True)

    server = BoundedThreadedServer((PROXY_BIND, PROXY_PORT), RobustHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")
        server.shutdown()