#!/usr/bin/env python3
"""
Self-hosted IPTV: serves M3U playlists that schedule video files from /movies and /series
with randomized order so the same show never appears twice in a row.
"""

import json
import random
import subprocess
import tempfile
import time
import xml.etree.ElementTree as ET
from collections import defaultdict
from datetime import datetime, timezone, timedelta
from pathlib import Path
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs, unquote
import mimetypes

# Per-channel schedule cache: channel_id -> (paths, durations). TTL 24h so EPG and live stay in sync.
_schedule_cache: dict[str, tuple[list[Path], list[float], float]] = {}
SCHEDULE_CACHE_TTL_SEC = 24 * 3600
EPG_DAYS = 14
DEFAULT_DURATION_SEC = 3600

MOVIES_ROOT = Path("/movies")
SERIES_ROOT = Path("/series")
DATA_ROOT = Path("/data")
CHANNELS_FILE = DATA_ROOT / "channels.json"
VIDEO_EXTENSIONS = {".mkv", ".mp4", ".avi", ".mov", ".m4v", ".webm", ".wmv"}


def get_channels():
    """Load channels from /data/channels.json."""
    if not CHANNELS_FILE.exists():
        return []
    with open(CHANNELS_FILE, encoding="utf-8") as f:
        data = json.load(f)
    return data.get("channels", [])


def resolve_path(path_key: str) -> Path | None:
    """Resolve a channel path key to an absolute Path. Returns None if invalid."""
    path_key = path_key.strip("/")
    if path_key == "movies":
        return MOVIES_ROOT
    if path_key.startswith("series/"):
        subpath = path_key[7:]  # len("series/")
        return SERIES_ROOT / subpath
    return None


def collect_videos(root: Path) -> list[tuple[str, Path]]:
    """Collect all video files under root. Returns list of (path_key, Path)."""
    if not root.exists() or not root.is_dir():
        return []
    videos = []
    try:
        for p in root.rglob("*"):
            if p.is_file() and p.suffix.lower() in VIDEO_EXTENSIONS:
                videos.append(p)
    except OSError:
        pass
    return videos


def collect_channel_videos(path_keys: list[str]) -> list[tuple[str, Path]]:
    """
    For each path key (e.g. 'movies', 'series/Breaking Bad'), collect videos.
    Returns list of (path_key, absolute_path) so we can group by path_key (show).
    """
    out = []
    for key in path_keys:
        root = resolve_path(key)
        if root is None:
            continue
        for vid in collect_videos(root):
            out.append((key, vid))
    return out


def shuffle_no_adjacent_same_show(items: list[tuple[str, Path]]) -> list[Path]:
    """
    Shuffle so that the same path_key (show) never appears twice in a row.
    Groups by path_key, shuffles each group, then interleaves round-robin.
    """
    if not items:
        return []
    by_key = defaultdict(list)
    for key, path in items:
        by_key[key].append(path)
    for key in by_key:
        random.shuffle(by_key[key])
    # Round-robin interleave so no two from same key are adjacent
    groups = list(by_key.values())
    result = []
    n = max(len(g) for g in groups)
    for i in range(n):
        for g in groups:
            if i < len(g):
                result.append(g[i])
    return result


def build_channel_path_list(path_keys: list[str]) -> list[Path]:
    """Build randomized list of absolute video paths for a channel (no same show adjacent)."""
    items = collect_channel_videos(path_keys)
    return shuffle_no_adjacent_same_show(items)


def get_duration(path: Path) -> float:
    """Return duration in seconds via ffprobe. Uses DEFAULT_DURATION_SEC on failure."""
    try:
        out = subprocess.run(
            [
                "ffprobe",
                "-v", "error",
                "-show_entries", "format=duration",
                "-of", "default=noprint_wrappers=1:nokey=1",
                str(path),
            ],
            capture_output=True,
            text=True,
            timeout=30,
        )
        if out.returncode == 0 and out.stdout.strip():
            return float(out.stdout.strip())
    except (subprocess.TimeoutExpired, ValueError, FileNotFoundError):
        pass
    return DEFAULT_DURATION_SEC


def get_or_build_schedule(channel_id: str) -> tuple[list[Path], list[float]] | None:
    """
    Return (paths, durations) for the channel, so live stream and EPG use the same schedule.
    Cached for SCHEDULE_CACHE_TTL_SEC.
    """
    now = time.time()
    if channel_id in _schedule_cache:
        paths, durations, ts = _schedule_cache[channel_id]
        if now - ts < SCHEDULE_CACHE_TTL_SEC:
            return paths, durations
    channels = get_channels()
    channel = next((c for c in channels if c.get("id") == channel_id), None)
    if not channel:
        return None
    paths = build_channel_path_list(channel.get("paths", []))
    if not paths:
        return None
    durations = [get_duration(p) for p in paths]
    _schedule_cache[channel_id] = (paths, durations, now)
    return paths, durations


def build_playlist(path_keys: list[str], base_url: str) -> list[tuple[str, str]]:
    """
    Build playlist as list of (display_name, stream_url).
    base_url is the base URL for stream links (e.g. http://host:port).
    """
    paths = build_channel_path_list(path_keys)
    playlist = []
    for p in paths:
        try:
            if MOVIES_ROOT in p.parents:
                rel = p.relative_to(MOVIES_ROOT)
                path_str = "movies/" + str(rel).replace("\\", "/")
            else:
                rel = p.relative_to(SERIES_ROOT)
                path_str = "series/" + str(rel).replace("\\", "/")
        except ValueError:
            path_str = "movies/" + p.name
        stream_url = f"{base_url}/stream?path={path_str}"
        display = p.stem
        playlist.append((display, stream_url))
    return playlist


def path_to_absolute(path_param: str) -> Path | None:
    """
    Convert a path query parameter to an absolute Path under /movies or /series.
    Path must start with 'movies/' or 'series/'. Prevents path traversal.
    """
    path_param = unquote(path_param).strip("/")
    if ".." in path_param or path_param.startswith("/"):
        return None
    if path_param.startswith("movies/"):
        sub = path_param[7:]
        root = MOVIES_ROOT
    elif path_param.startswith("series/"):
        sub = path_param[7:]
        root = SERIES_ROOT
    else:
        return None
    candidate = (root / sub).resolve()
    try:
        candidate.relative_to(root)
    except ValueError:
        return None
    if candidate.exists() and candidate.is_file() and candidate.suffix.lower() in VIDEO_EXTENSIONS:
        return candidate
    return None


def get_base_url(handler: BaseHTTPRequestHandler) -> str:
    """Build base URL for playlist links from request."""
    host = handler.headers.get("Host", "localhost:8080")
    return f"http://{host}"


def _format_xmltv_time(dt: datetime) -> str:
    """Format as XMLTV: YYYYMMDDHHmmss +0000"""
    tz = dt.strftime("%z") if dt.tzinfo else "+0000"
    return dt.strftime("%Y%m%d%H%M%S ") + tz


def build_epg_xml() -> str:
    """
    Build XMLTV EPG for all channels. Schedule repeats from midnight UTC;
    programmes generated for EPG_DAYS. Uses cached schedule per channel.
    """
    channels = get_channels()
    if not channels:
        return '<?xml version="1.0" encoding="UTF-8"?>\n<tv></tv>'

    now = datetime.now(timezone.utc)
    epoch = now.replace(hour=0, minute=0, second=0, microsecond=0)
    end_epoch = epoch + timedelta(days=EPG_DAYS)

    root = ET.Element("tv")
    root.set("source-info-name", "self-hosted-iptv")
    root.set("generator-info-name", "self-hosted-iptv")

    for ch in channels:
        cid = ch.get("id", ch.get("name", ""))
        name = ch.get("name", cid)
        chan_el = ET.SubElement(root, "channel", id=cid)
        ET.SubElement(chan_el, "display-name").text = name

    for ch in channels:
        cid = ch.get("id", ch.get("name", ""))
        schedule = get_or_build_schedule(cid)
        if not schedule:
            continue
        paths, durations = schedule
        cycle_duration_sec = sum(durations)
        if cycle_duration_sec <= 0:
            continue
        cycle_delta = timedelta(seconds=cycle_duration_sec)
        offsets_sec = [0.0]
        for d in durations[:-1]:
            offsets_sec.append(offsets_sec[-1] + d)

        cycle_start = epoch
        while cycle_start < end_epoch:
            for i, (path, dur) in enumerate(zip(paths, durations)):
                start_dt = cycle_start + timedelta(seconds=offsets_sec[i])
                stop_dt = start_dt + timedelta(seconds=dur)
                if stop_dt <= now:
                    continue
                if start_dt >= end_epoch:
                    break
                prog = ET.SubElement(root, "programme",
                    start=_format_xmltv_time(start_dt),
                    stop=_format_xmltv_time(stop_dt),
                    channel=cid,
                )
                title = path.stem
                ET.SubElement(prog, "title", lang="en").text = title
            cycle_start += cycle_delta

    ET.indent(root, space="  ")
    return '<?xml version="1.0" encoding="UTF-8"?>\n' + ET.tostring(root, encoding="unicode", default_namespace=None)


def write_concat_list(paths: list[Path], fd) -> None:
    """Write FFmpeg concat demuxer list to file. Escapes paths for safety."""
    for p in paths:
        # Concat format: file 'path' — escape ' as '\''
        path_str = str(p.resolve()).replace("\\", "\\\\").replace("'", "'\\''")
        fd.write(f"file '{path_str}'\n")
    fd.flush()


def stream_live_channel(paths: list[Path], wfile) -> None:
    """
    Run FFmpeg to output a continuous MPEG-TS stream from paths, looping forever.
    Pipe output to wfile. Stops when wfile write fails (client disconnect).
    """
    if not paths:
        return
    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
        write_concat_list(paths, f)
        list_path = f.name
    try:
        proc = subprocess.Popen(
            [
                "ffmpeg",
                "-stream_loop", "-1",
                "-f", "concat", "-safe", "0", "-i", list_path,
                "-c", "copy",
                "-f", "mpegts",
                "-",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            start_new_session=True,
        )
        try:
            while True:
                chunk = proc.stdout.read(65536)
                if not chunk:
                    break
                wfile.write(chunk)
                wfile.flush()
        except (BrokenPipeError, ConnectionResetError, OSError):
            pass
        finally:
            proc.terminate()
            try:
                proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                proc.kill()
    finally:
        Path(list_path).unlink(missing_ok=True)


class IPTVHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        pass  # quiet by default; set to super().log_message for debugging

    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path.rstrip("/") or "/"
        qs = parse_qs(parsed.query)

        if path == "/" or path == "/index.html":
            self.send_index()
            return
        if path == "/playlist.m3u" or path == "/channels.m3u":
            self.send_channels_m3u()
            return
        if path.startswith("/channel/") and path.endswith("/playlist.m3u"):
            channel_id = path.split("/")[2]
            self.send_channel_m3u(channel_id)
            return
        if path == "/stream":
            path_param = qs.get("path", [""])[0]
            self.send_stream(path_param)
            return
        if path.startswith("/live/"):
            parts = path.split("/")
            if len(parts) == 3 and parts[1] == "live":
                self.send_live_stream(parts[2])
                return
        if path == "/epg.xml" or path == "/xmltv.xml":
            self.send_epg()
            return
        self.send_error(404, "Not found")

    def send_index(self):
        channels = get_channels()
        body = "<!DOCTYPE html><html><head><meta charset='utf-8'><title>Self-hosted IPTV</title></head><body>"
        body += "<h1>Self-hosted IPTV</h1><p>Playlist (M3U):</p>"
        body += f"<p><code>{get_base_url(self)}/playlist.m3u</code></p>"
        body += "<p>EPG (XMLTV, for Plex DVR):</p>"
        body += f"<p><code>{get_base_url(self)}/epg.xml</code></p>"
        body += "<h2>Channels</h2><ul>"
        for ch in channels:
            cid = ch.get("id", ch.get("name", ""))
            body += f"<li><a href='/channel/{cid}/playlist.m3u'>{ch.get('name', cid)}</a></li>"
        body += "</ul></body></html>"
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body.encode("utf-8"))))
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def send_channels_m3u(self):
        """Master M3U for Plex DVR: one live stream URL per channel."""
        base = get_base_url(self)
        channels = get_channels()
        lines = ["#EXTM3U"]
        for ch in channels:
            cid = ch.get("id", ch.get("name", ""))
            name = ch.get("name", cid)
            lines.append(f'#EXTINF:-1,{name}')
            lines.append(f"{base}/live/{cid}")
        body = "\n".join(lines) + "\n"
        self.send_response(200)
        self.send_header("Content-Type", "application/x-mpegURL; charset=utf-8")
        self.send_header("Content-Length", str(len(body.encode("utf-8"))))
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def send_channel_m3u(self, channel_id: str):
        """Single channel M3U with randomized video entries (no same show adjacent)."""
        channels = get_channels()
        channel = next((c for c in channels if c.get("id") == channel_id), None)
        if not channel:
            self.send_error(404, "Channel not found")
            return
        path_keys = channel.get("paths", [])
        base = get_base_url(self)
        playlist = build_playlist(path_keys, base)
        lines = ["#EXTM3U"]
        for name, url in playlist:
            lines.append(f"#EXTINF:-1,{name}")
            lines.append(url)
        body = "\n".join(lines) + "\n"
        self.send_response(200)
        self.send_header("Content-Type", "application/x-mpegURL; charset=utf-8")
        self.send_header("Content-Length", str(len(body.encode("utf-8"))))
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def send_epg(self):
        """Serve XMLTV EPG (programme schedule for Plex DVR)."""
        body = build_epg_xml()
        self.send_response(200)
        self.send_header("Content-Type", "application/xml; charset=utf-8")
        self.send_header("Content-Length", str(len(body.encode("utf-8"))))
        self.send_header("Cache-Control", "public, max-age=300")
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def send_live_stream(self, channel_id: str):
        """Stream channel as continuous MPEG-TS (for Plex DVR). Uses cached schedule so EPG matches."""
        schedule = get_or_build_schedule(channel_id)
        if not schedule:
            self.send_error(404, "Channel not found or no videos")
            return
        paths, _ = schedule
        if not paths:
            self.send_error(404, "No videos in channel")
            return
        self.send_response(200)
        self.send_header("Content-Type", "video/MP2T")
        self.send_header("Cache-Control", "no-cache, no-store")
        self.end_headers()
        stream_live_channel(paths, self.wfile)

    def send_stream(self, path_param: str):
        """Stream a video file. path is relative to /movies or /series."""
        abs_path = path_to_absolute(path_param)
        if abs_path is None:
            self.send_error(404, "File not found")
            return
        content_type, _ = mimetypes.guess_type(str(abs_path))
        content_type = content_type or "video/mp4"
        size = abs_path.stat().st_size
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(size))
        self.send_header("Accept-Ranges", "bytes")
        self.end_headers()
        with open(abs_path, "rb") as f:
            self.wfile.write(f.read())


def main():
    if not CHANNELS_FILE.exists():
        print(f"Warning: {CHANNELS_FILE} not found. Create it from data/channels.json.example")
    port = 8080
    server = HTTPServer(("0.0.0.0", port), IPTVHandler)
    print(f"Serving at http://0.0.0.0:{port}/  (playlist: http://<host>:{port}/playlist.m3u)")
    server.serve_forever()


if __name__ == "__main__":
    main()