import io
import json
import os
import struct
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from host.common import safe_filename, select_streams, validate_download, youtube_url
from host.media import encoding_plan, ffmpeg_command
from host.protocol import MAX_MESSAGE, read_message, write_message
from host.worker import publish_file
from host.main import NativeHost
from host.session import apply_cookies, validate_cookies, public_error


class HostTests(unittest.TestCase):
    def test_fragment_access_failure_retries_and_preserves_the_original_error(self):
        from unittest.mock import MagicMock
        from yt_dlp.utils import DownloadError
        from host.worker import download
        first, second = MagicMock(), MagicMock()
        first.__enter__.return_value = first; second.__enter__.return_value = second
        handles = []
        def fail(ydl):
            # Reproduce a downloader retaining its .part handle during failure.
            part = Path(ydl.params["outtmpl"].replace("%(ext)s", "mp4") + ".part")
            handles.append(part.open("wb"))
            raise DownloadError("HTTP Error 403: Forbidden")
        def create_ydl(params):
            current = first if not handles else second
            current.params = params
            current.extract_info.side_effect = lambda *args, **kwargs: fail(current)
            return current
        with tempfile.TemporaryDirectory() as folder:
            try:
                with patch("yt_dlp.YoutubeDL", side_effect=create_ydl), patch("host.worker.ydl_options", return_value={}), \
                     patch("host.worker.time.sleep"), patch("host.worker.emit") as emit:
                    with self.assertRaisesRegex(DownloadError, "HTTP Error 403"):
                        download({"url": "https://youtu.be/jNQXAC9IVRw", "folder": folder}, [])
                    self.assertEqual(len(handles), 2)
                    self.assertTrue(any("повторяю загрузку" in str(call) for call in emit.call_args_list))
            finally:
                for handle in handles: handle.close()

    def test_transient_youtube_challenge_is_retried_once_without_browser_cookies(self):
        from unittest.mock import MagicMock
        from yt_dlp.utils import DownloadError
        first, second = MagicMock(), MagicMock()
        first.__enter__.return_value = first; second.__enter__.return_value = second
        first.extract_info.side_effect = DownloadError("Sign in to confirm you’re not a bot")
        second.extract_info.return_value = {"id": "jNQXAC9IVRw"}
        with patch("yt_dlp.YoutubeDL", side_effect=[first, second]), patch("host.worker.ydl_options", return_value={}), \
             patch("host.worker.summarize", side_effect=lambda info: info), patch("host.worker.time.sleep"):
            from host.worker import inspect
            self.assertEqual(inspect("https://youtu.be/jNQXAC9IVRw")["id"], "jNQXAC9IVRw")
        first.cookiejar.set_cookie.assert_not_called(); second.cookiejar.set_cookie.assert_not_called()

    def test_browser_session_is_youtube_only_and_never_saved_to_a_cookie_file(self):
        from urllib.request import Request
        from yt_dlp import YoutubeDL
        from yt_dlp.cookies import YoutubeDLCookieJar
        cookie = {"name": "SAPISID", "value": "synthetic-cookie", "domain": ".youtube.com", "path": "/", "secure": True, "httpOnly": True}
        with patch.object(YoutubeDLCookieJar, "save", side_effect=AssertionError("Cookies must stay in memory")):
            with YoutubeDL({"quiet": True}) as ydl:
                apply_cookies(ydl, [cookie])
                self.assertIsNone(ydl.cookiejar.filename)
                request = Request("https://www.youtube.com/watch?v=jNQXAC9IVRw")
                ydl.cookiejar.add_cookie_header(request)
                self.assertIn("SAPISID=synthetic-cookie", request.get_header("Cookie"))
                unrelated = Request("https://www.google.com/")
                ydl.cookiejar.add_cookie_header(unrelated)
                self.assertIsNone(unrelated.get_header("Cookie"))
        for invalid in [{**cookie, "domain": ".google.com"}, {**cookie, "domain": "youtube.com.evil.test"},
                        {**cookie, "value": "secret\r\nInjected: header"}, {**cookie, "expirationDate": float("nan")}, {**cookie, "name": "bad=name"}]:
            with self.assertRaises(ValueError): validate_cookies([invalid])

    def test_native_inspection_passes_cookies_only_to_worker_and_not_to_reply(self):
        with tempfile.TemporaryDirectory() as folder, patch("host.main.ROOT", Path(folder)):
            output = io.BytesIO(); host = NativeHost(output)
            cookie = {"name": "SID", "value": "synthetic-session", "domain": ".youtube.com", "path": "/"}
            try:
                with patch.object(host, "worker", return_value={"title": "Test"}) as worker:
                    host.handle({"id": "inspect-auth", "action": "inspect", "payload": {"url": "https://youtu.be/jNQXAC9IVRw", "browserCookies": [cookie]}})
                    self.assertEqual(worker.call_args.args[1]["browserCookies"][0]["value"], cookie["value"])
                self.assertNotIn(b"synthetic-session", output.getvalue())
                self.assertNotIn("browserCookies", json.dumps(host.jobs))
            finally: host.close()

    def test_youtube_login_error_is_distinct_from_helper_connection_failure(self):
        message = public_error(ValueError("ERROR: Sign in to confirm you’re not a bot"))
        self.assertIn("Использовать вход в YouTube", message)
        self.assertIn("Помощник работает", message)
        self.assertNotIn("cookies-from-browser", message)
        self.assertIn("отклонил вход", public_error(ValueError("Sign in to confirm"), True))

    def test_youtube_urls_are_canonical_and_do_not_allow_arbitrary_hosts(self):
        expected = "https://www.youtube.com/watch?v=jNQXAC9IVRw"
        for value in [expected + "&list=other", "https://youtu.be/jNQXAC9IVRw?t=2", "https://www.youtube.com/shorts/jNQXAC9IVRw"]:
            self.assertEqual(youtube_url(value), expected)
        for value in ["https://youtube.com.evil.test/watch?v=jNQXAC9IVRw", "https://127.0.0.1/watch?v=jNQXAC9IVRw", "https://youtube.com@evil.test/watch?v=jNQXAC9IVRw", "file:///test", "https://www.youtube.com:444/watch?v=jNQXAC9IVRw", "https://www.youtube.com/@channel"]:
            with self.assertRaises(ValueError): youtube_url(value)

    def test_separate_audio_is_added_and_video_only_stays_single(self):
        info = {"formats": [{"format_id": "18", "vcodec": "h264", "acodec": "aac", "height": 360},
                            {"format_id": "137", "vcodec": "h264", "acodec": "none", "height": 1080},
                            {"format_id": "140", "vcodec": "none", "acodec": "aac"}]}
        base = {"mode": "video", "videoId": "137", "audioId": "best"}
        self.assertEqual(select_streams(info, base), "137+140")
        self.assertEqual(select_streams(info, {**base, "mode": "video_only"}), "137")
        self.assertEqual(select_streams(info, {**base, "videoId": "18"}), "18")
        self.assertEqual(select_streams(info, {**base, "mode": "audio"}), "140")

    def test_download_does_not_accept_format_selector_injection(self):
        with tempfile.TemporaryDirectory() as folder:
            base = {"url": "https://youtu.be/jNQXAC9IVRw", "folder": folder, "videoId": "137+bestaudio"}
            with self.assertRaises(ValueError): validate_download(base)
            self.assertEqual(validate_download({**base, "videoId": "137"})["container"], "mp4")
            with self.assertRaises(ValueError): validate_download({**base, "videoId": "137", "container": "exe"})

    def test_gpu_is_used_only_when_video_needs_transcoding(self):
        media = {"streams": [{"codec_type": "video", "codec_name": "vp9", "pix_fmt": "yuv420p"}, {"codec_type": "audio", "codec_name": "opus"}]}
        options = {"mode": "video", "container": "mp4", "codec": "auto", "engine": "auto"}
        plan = encoding_plan(options, media, ["h264", "hevc", "av1"])
        self.assertEqual((plan["video"], plan["audio"], plan["engine"]), ("h264", "aac", "nvenc"))
        copied = encoding_plan({**options, "container": "mkv"}, media, ["h264"])
        self.assertEqual((copied["video"], copied["audio"], copied["engine"]), ("copy", "copy", "copy"))
        cpu = encoding_plan(options, media, [])
        self.assertEqual(cpu["engine"], "cpu")
        with self.assertRaises(ValueError): encoding_plan({**options, "engine": "nvenc"}, media, [])
        with self.assertRaises(ValueError): encoding_plan({**options, "codec": "copy"}, media, [])

    def test_hdr_is_not_silently_converted_to_h264(self):
        media = {"streams": [{"codec_type": "video", "codec_name": "vp9", "pix_fmt": "yuv420p10le", "color_transfer": "smpte2084"}, {"codec_type": "audio", "codec_name": "opus"}]}
        options = {"mode": "video", "container": "mp4", "codec": "auto", "engine": "auto"}
        self.assertEqual(encoding_plan(options, media, ["hevc"])["video"], "hevc")
        with self.assertRaises(ValueError): encoding_plan({**options, "codec": "h264"}, media, ["h264"])

    def test_amf_requires_measured_codec_and_10bit_support(self):
        media = {"streams": [{"codec_type": "video", "codec_name": "vp9", "pix_fmt": "yuv420p"}, {"codec_type": "audio", "codec_name": "opus"}]}
        options = {"mode": "video", "container": "mp4", "codec": "hevc", "engine": "auto", "preset": "balanced", "bitrate": 256}
        self.assertEqual(encoding_plan(options, media, {"amf": ["hevc"]})["engine"], "amf")
        self.assertEqual(encoding_plan(options, media, {"nvenc": ["hevc"], "amf": ["hevc"]})["engine"], "nvenc")
        with self.assertRaisesRegex(ValueError, "AMD AMF"):
            encoding_plan({**options, "engine": "amf"}, media, {})
        media["streams"][0].update(pix_fmt="yuv420p10le", color_transfer="smpte2084")
        self.assertEqual(encoding_plan(options, media, {"amf": ["hevc"]})["engine"], "cpu")
        plan = encoding_plan({**options, "engine": "amf"}, media, {"amf": ["hevc"], "amf10bit": ["hevc"]})
        self.assertTrue(plan["tenBit"])
        with patch("host.media.tool", return_value="ffmpeg"):
            command = ffmpeg_command("source.mkv", "out.mp4", options, plan)
        self.assertIn("hevc_amf", command); self.assertIn("p010le", command); self.assertIn("main10", command)
        self.assertIn("d3d11va", command); self.assertNotIn("cuda", command)

    def test_audio_and_video_only_output_maps(self):
        source = {"streams": [{"codec_type": "video", "codec_name": "h264"}, {"codec_type": "audio", "codec_name": "aac"}]}
        options = {"mode": "video_only", "container": "mp4", "codec": "auto", "engine": "auto", "preset": "balanced", "bitrate": 256}
        with patch("host.media.tool", return_value="ffmpeg"):
            command = ffmpeg_command("source.mkv", "out.mp4", options, encoding_plan(options, source, []))
        self.assertIn("-an", command)
        self.assertNotIn("0:a:0", command)
        audio = encoding_plan({**options, "mode": "audio", "container": "mp3"}, source, [])
        self.assertEqual((audio["video"], audio["audio"], audio["engine"]), (None, "libmp3lame", "cpu"))

    def test_protocol_uses_utf8_byte_lengths_and_handles_partial_reads(self):
        message = {"text": "Русский путь 🎬"}
        stream = io.BytesIO()
        write_message(stream, message)
        self.assertEqual(struct.unpack("<I", stream.getvalue()[:4])[0], len(stream.getvalue()) - 4)
        class Partial(io.BytesIO):
            def read(self, count=-1): return super().read(min(count, 2))
        self.assertEqual(read_message(Partial(stream.getvalue())), message)
        self.assertIsNone(read_message(io.BytesIO()))
        with self.assertRaises(ValueError): read_message(io.BytesIO(struct.pack("<I", MAX_MESSAGE + 1)))
        with self.assertRaises(EOFError): read_message(io.BytesIO(struct.pack("<I", 5) + b"{}"))

    def test_filename_is_safe_and_existing_files_are_never_overwritten(self):
        self.assertEqual(safe_filename("../CON:bad/name", "video"), "_CON_bad_name")
        self.assertEqual(safe_filename("CON", "video"), "_CON")
        with tempfile.TemporaryDirectory() as folder:
            root = Path(folder); original = root / "video.mp4"; original.write_bytes(b"original")
            source = root / "temp.mp4"; source.write_bytes(b"new")
            saved = publish_file(source, root, "video", "mp4")
            self.assertEqual(saved.name, "video (1).mp4")
            self.assertEqual(original.read_bytes(), b"original")
            self.assertEqual(saved.read_bytes(), b"new")

    def test_dependency_install_refreshes_capabilities_before_reporting_completion(self):
        with tempfile.TemporaryDirectory() as folder, patch("host.main.ROOT", Path(folder)):
            host = NativeHost(io.BytesIO())
            host._capabilities = {"ready": False}
            host.jobs["install"] = {"id": "install", "status": "running", "kind": "dependency"}
            try:
                with patch.object(host, "worker", return_value={"title": "Node.js"}), patch.object(host, "event") as event:
                    event.side_effect = lambda *args, **kwargs: self.assertIsNone(host._capabilities)
                    host.run_job("install", {"dependency": "node"}, "dependency_install")
                    self.assertEqual(event.call_args.kwargs["status"], "complete")
                with patch.object(host, "worker", side_effect=ValueError("install failed")), patch.object(host, "event") as event:
                    host._capabilities = {"ready": False}
                    event.side_effect = lambda *args, **kwargs: self.assertIsNone(host._capabilities)
                    host.run_job("install", {"dependency": "node"}, "dependency_install")
                    self.assertEqual(event.call_args.kwargs["status"], "error")
            finally:
                host.close()

    def test_completed_file_can_be_revealed_after_native_host_restart(self):
        with tempfile.TemporaryDirectory() as folder, patch("host.main.ROOT", Path(folder)):
            saved = Path(folder) / "finished.mp4"
            saved.write_bytes(b"media")
            first = NativeHost(io.BytesIO())
            first.jobs["completed"] = {"id": "completed", "status": "running", "createdAt": 1}
            first.event("completed", status="complete", result={"path": str(saved)})
            first.close()
            output = io.BytesIO()
            restarted = NativeHost(output)
            try:
                with patch("host.main.subprocess.Popen") as launch:
                    restarted.handle({"id": "reveal", "action": "reveal", "payload": {"jobId": "completed"}})
                output.seek(0)
                self.assertTrue(read_message(output)["ok"])
                if os.name == "nt":
                    self.assertEqual(launch.call_args.args[0][-1], str(saved))
            finally:
                restarted.close()


if __name__ == "__main__": unittest.main()