from __future__ import annotations

import contextlib
import functools
import http.server
from pathlib import Path
import shutil
import subprocess
import threading

from n_m3u8dl_py.cli import main
from n_m3u8dl_py.hls import HLSExtractor
from n_m3u8dl_py.http import HttpClient


def test_hls_master_and_media_playlist(tmp_path: Path) -> None:
    (tmp_path / "video.m3u8").write_text("#EXTM3U\n#EXTINF:2,\na.ts\n#EXT-X-ENDLIST\n")
    master = """#EXTM3U
#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID="aud",NAME="English",LANGUAGE="en",URI="audio.m3u8"
#EXT-X-STREAM-INF:BANDWIDTH=1000,RESOLUTION=640x360,AUDIO="aud"
video.m3u8
"""
    extractor = HLSExtractor((tmp_path / "master.m3u8").as_uri(), HttpClient())
    streams = extractor.extract(master)
    assert len(streams) == 2
    video = next(stream for stream in streams if stream.resolution == "640x360")
    extractor.fetch_playlists([video])
    assert video.playlist is not None
    assert video.playlist.segments[0].url.endswith("/a.ts")


def test_local_hls_download_and_binary_merge(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"first")
    (tmp_path / "two.ts").write_bytes(b"second")
    manifest = tmp_path / "test.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXTINF:1,\ntwo.ts\n#EXT-X-ENDLIST\n")
    output = tmp_path / "out"
    assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "joined", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
    assert (output / "joined.ts").read_bytes() == b"firstsecond"


def test_local_hls_ffmpeg_merge_remuxes_ts_to_mp4(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"first")
    manifest = tmp_path / "test.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
    ffmpeg = tmp_path / "fake-ffmpeg"
    ffmpeg.write_text("#!/bin/sh\nfor argument do output=$argument; done\nprintf remuxed > \"$output\"\n", encoding="utf-8")
    ffmpeg.chmod(0o755)
    output = tmp_path / "out"

    assert main([
        str(manifest),
        "--auto-select",
        "--ffmpeg-binary-path",
        str(ffmpeg),
        "--save-name",
        "episode",
        "--save-dir",
        str(output),
        "--tmp-dir",
        str(tmp_path / "tmp"),
    ]) == 0

    assert (output / "episode.mp4").read_bytes() == b"remuxed"
    assert not (output / "episode.ts").exists()


def test_local_hls_save_pattern(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"first")
    manifest = tmp_path / "test.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
    output = tmp_path / "out"
    assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "episode", "--save-pattern", "<SaveName>_<MediaType>.<Ext>", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
    assert (output / "episode_video.ts").read_bytes() == b"first"


def test_local_hls_file_uri_decodes_spaces(tmp_path: Path) -> None:
    source = tmp_path / "directory with spaces"
    source.mkdir()
    (source / "one.ts").write_bytes(b"space-safe")
    manifest = source / "test.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
    output = tmp_path / "out"
    assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "spaced", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
    assert (output / "spaced.ts").read_bytes() == b"space-safe"


def test_http_hls_download(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"served-over-http")
    (tmp_path / "test.m3u8").write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
    output = tmp_path / "out"
    with _http_server(tmp_path) as base_url:
        assert main([f"{base_url}/test.m3u8", "--no-use-system-proxy", "--auto-select", "--binary-merge", "--save-name", "http", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
    assert (output / "http.ts").read_bytes() == b"served-over-http"


def test_local_hls_aes128_download(tmp_path: Path) -> None:
    key = bytes(range(16))
    iv = bytes(16)
    plaintext = b"encrypted payload"
    padding = 16 - len(plaintext) % 16
    ciphertext = subprocess.run(
        [shutil.which("openssl") or "openssl", "enc", "-aes-128-cbc", "-K", key.hex(), "-iv", iv.hex(), "-nopad"],
        input=plaintext + bytes([padding]) * padding,
        capture_output=True,
        check=True,
    ).stdout
    (tmp_path / "key.bin").write_bytes(key)
    (tmp_path / "encrypted.ts").write_bytes(ciphertext)
    manifest = tmp_path / "encrypted.m3u8"
    manifest.write_text("#EXTM3U\n#EXT-X-KEY:METHOD=AES-128,URI=\"key.bin\",IV=0x00000000000000000000000000000000\n#EXTINF:1,\nencrypted.ts\n#EXT-X-ENDLIST\n")
    output = tmp_path / "out"
    assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "plain", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
    assert (output / "plain.ts").read_bytes() == plaintext


@contextlib.contextmanager
def _http_server(directory: Path):
    handler = functools.partial(http.server.SimpleHTTPRequestHandler, directory=str(directory))
    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), handler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        yield f"http://127.0.0.1:{server.server_port}"
    finally:
        server.shutdown()
        thread.join()