# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for UnifiedMobileController video recording and playback features."""

import subprocess
import time
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4

import pytest
import cv2

from artemis.context import ArtemisContext
from artemis.controllers.unified_controller import UnifiedMobileController
from artemis.drivers.mock.mock_driver import MockDeviceDriver
from artemis.utils.video import (
    RecordingSession,
    build_scrcpy_record_command,
    get_ffmpeg_path,
    get_active_session,
    normalize_recording_to_mp4,
    render_timeline_clip,
    remove_active_session,
    set_active_session,
)


def test_scrcpy_recording_locks_each_segment_orientation(tmp_path):
    output_path = tmp_path / "recording.mkv"

    command = build_scrcpy_record_command("scrcpy", "device-1", output_path)

    assert "--capture-orientation=@" in command
    assert command[command.index("--record") + 1] == str(output_path)


@pytest.mark.asyncio
async def test_normalize_recording_handles_resolution_change(tmp_path):
    """A portrait/landscape H.264 track must become one fixed-size MP4."""
    ffmpeg = get_ffmpeg_path()
    portrait = tmp_path / "portrait.mkv"
    landscape = tmp_path / "landscape.mkv"
    dynamic = tmp_path / "dynamic.mkv"
    output = tmp_path / "recording.mp4"

    for size, path in (("108x242", portrait), ("242x108", landscape)):
        subprocess.run(
            [
                ffmpeg,
                "-y",
                "-f",
                "lavfi",
                "-i",
                f"testsrc=size={size}:rate=10:duration=0.5",
                "-c:v",
                "libx264",
                "-pix_fmt",
                "yuv420p",
                str(path),
            ],
            check=True,
            capture_output=True,
        )

    concat_file = tmp_path / "segments.txt"
    concat_file.write_text(
        f"file '{portrait.as_posix()}'\nfile '{landscape.as_posix()}'\n",
        encoding="utf-8",
    )
    subprocess.run(
        [
            ffmpeg,
            "-y",
            "-f",
            "concat",
            "-safe",
            "0",
            "-i",
            str(concat_file),
            "-c",
            "copy",
            str(dynamic),
        ],
        check=True,
        capture_output=True,
    )

    assert await normalize_recording_to_mp4(dynamic, output, 108, 242)

    capture = cv2.VideoCapture(str(output))
    dimensions = set()
    frame_count = 0
    while True:
        ok, frame = capture.read()
        if not ok:
            break
        frame_count += 1
        dimensions.add((frame.shape[1], frame.shape[0]))
    capture.release()

    assert frame_count >= 8
    assert dimensions == {(108, 242)}


@pytest.mark.asyncio
async def test_analyzer_clip_stays_continuous_across_orientation_segments(tmp_path):
    """The analyzer must still receive one decodable MP4 across a rotation."""
    ffmpeg = get_ffmpeg_path()
    portrait = tmp_path / "recording.mkv"
    landscape = tmp_path / "recording_001.mkv"
    output = tmp_path / "agent_clip.mp4"
    for color, size, path in (
        ("red", "108x242", portrait),
        ("blue", "242x108", landscape),
    ):
        subprocess.run(
            [
                ffmpeg,
                "-y",
                "-f",
                "lavfi",
                "-i",
                f"color={color}:size={size}:rate=15:duration=1",
                "-c:v",
                "libx264",
                "-pix_fmt",
                "yuv420p",
                str(path),
            ],
            check=True,
            capture_output=True,
        )

    assert await render_timeline_clip(
        [
            {"path": portrait, "start": 0.0, "end": 1.0},
            {"path": landscape, "start": 1.0, "end": 2.0},
        ],
        0.25,
        1.75,
        output,
        canvas_width=180,
        canvas_height=320,
    )

    capture = cv2.VideoCapture(str(output))
    dimensions = set()
    frames = 0
    while True:
        ok, frame = capture.read()
        if not ok:
            break
        dimensions.add((frame.shape[1], frame.shape[0]))
        frames += 1
    capture.release()
    assert dimensions == {(180, 320)}
    assert frames >= 20


@pytest.fixture
def mock_ctx(tmp_path):
    ctx = MagicMock(spec=ArtemisContext)
    ctx.device = MagicMock()
    ctx.device.device_id = "emulator-5554"
    ctx.device.mobile_platform = "android"

    # DataEngine mock
    ctx.data_engine = MagicMock()
    ctx.data_engine.current_session_id = uuid4()
    ctx.data_engine.session_start_time = time.time()
    ctx.data_engine.storage = MagicMock()

    # Mock driver
    mock_driver = MockDeviceDriver(device_id="emulator-5554")
    ctx._active_driver = mock_driver
    return ctx


@pytest.mark.asyncio
async def test_unified_controller_start_recording(mock_ctx, tmp_path):
    controller = UnifiedMobileController(mock_ctx)
    remove_active_session("emulator-5554")

    # Mock scrcpy subprocess
    mock_proc = MagicMock()
    mock_proc.returncode = None

    with patch("asyncio.create_subprocess_exec", AsyncMock(return_value=mock_proc)):
        with patch(
            "artemis.controllers.unified_controller.get_android_display_state",
            AsyncMock(return_value=(0, 1080, 2424)),
        ):
            with patch("asyncio.sleep", AsyncMock()):
                res = await controller.start_video_recording(output_dir=tmp_path)

                assert res.success is True
                assert get_active_session("emulator-5554") is not None
                assert mock_ctx.data_engine.record_video_start.called

                # Calling start again should report already in progress
                res2 = await controller.start_video_recording(output_dir=tmp_path)
                assert res2.success is False
                assert "already in progress" in res2.message

    remove_active_session("emulator-5554")


@pytest.mark.asyncio
async def test_unified_controller_stop_recording(mock_ctx, tmp_path):
    controller = UnifiedMobileController(mock_ctx)
    remove_active_session("emulator-5554")

    # Create dummy recording.mkv
    mkv_path = tmp_path / "recording.mkv"
    mkv_path.write_bytes(b"dummy mkv content")

    mock_proc = MagicMock()
    mock_proc.returncode = None
    mock_proc.terminate = MagicMock()
    mock_proc.wait = AsyncMock()

    session = RecordingSession(
        video_id=uuid4(),
        device_id="emulator-5554",
        start_time=time.time() - 10.0,
        data_engine_start_time=time.time() - 10.0,
        local_video_path=mkv_path,
        process=mock_proc,
    )
    set_active_session("emulator-5554", session)

    with patch(
        "artemis.controllers.unified_controller.remux_recording_to_mp4",
        AsyncMock(side_effect=lambda src, dst: (dst.write_bytes(b"mp4 content"), True)[1]),
    ), patch(
        "artemis.controllers.unified_controller.write_recording_manifest",
        AsyncMock(return_value=tmp_path / "recording.json"),
    ):
        res = await controller.stop_video_recording()

        assert res.success is True
        assert res.video_path is not None
        assert str(res.video_path).endswith("recording.mp4")
        assert get_active_session("emulator-5554") is None
        assert mock_ctx.data_engine.record_video_stop.called


@pytest.mark.asyncio
async def test_unified_controller_extract_segment_metadata(mock_ctx, tmp_path):
    controller = UnifiedMobileController(mock_ctx)
    remove_active_session("emulator-5554")

    mkv_path = tmp_path / "recording.mkv"
    mkv_path.write_bytes(b"dummy mkv content")

    mock_proc = MagicMock()
    mock_proc.returncode = None

    session = RecordingSession(
        video_id=uuid4(),
        device_id="emulator-5554",
        start_time=time.time() - 20.0,
        data_engine_start_time=time.time() - 20.0,
        local_video_path=mkv_path,
        process=mock_proc,
    )
    set_active_session("emulator-5554", session)

    with patch(
        "artemis.controllers.unified_controller.render_timeline_clip",
        AsyncMock(side_effect=lambda src, s, e, dst: dst.write_bytes(b"segment mp4") or True),
    ):
        res = await controller.extract_segment_metadata(start_time=2.0, end_time=8.0)

        assert res.success is True
        assert res.video_path is not None
        assert res.video_path.exists()
        # Ensure session is still active (not stopped by extraction!)
        assert get_active_session("emulator-5554") is not None
    remove_active_session("emulator-5554")


def test_data_engine_video_lifecycle(tmp_path):
    from artemis.data_engine.engine import DataEngine
    from artemis.context import ArtemisContext

    mock_c = MagicMock(spec=ArtemisContext)
    mock_c.execution_setup = MagicMock(traces_path=str(tmp_path / "traces"))
    mock_c.device = None
    engine = DataEngine(mock_c)
    session_id = engine.start_session(goal="Test Video Engine")

    vid = uuid4()
    vpath = tmp_path / "recording.mp4"
    vpath.write_bytes(b"video bytes")

    # Start video recording
    engine.record_video_start(vid, "device-1", vpath)
    rec = engine.storage.get_video_recording(vid)
    assert rec is not None
    assert rec.device_id == "device-1"

    # Stop video recording
    engine.record_video_stop(vid, "device-1", vpath, time.time() - 10.0, time.time())
    sess = engine.storage.get_session(session_id)
    assert sess.video_filepath == str(vpath)

    # Move video path on finalize
    new_vpath = tmp_path / "archived" / "recording.mp4"
    new_vpath.parent.mkdir(parents=True, exist_ok=True)
    new_vpath.write_bytes(b"video bytes")
    engine.update_video_path(new_vpath)

    sess_updated = engine.storage.get_session(session_id)
    assert sess_updated.video_filepath == str(new_vpath)


@pytest.mark.asyncio
async def test_unified_controller_crash_recovery_and_multi_segment(mock_ctx, tmp_path):
    """Verify that if recording is interrupted, segments are auto-recovered and concatenated seamlessly."""
    controller = UnifiedMobileController(mock_ctx)
    remove_active_session("emulator-5554")

    seg1 = tmp_path / "recording_0.mkv"
    seg1.write_bytes(b"segment 0")
    seg2 = tmp_path / "recording_1.mkv"
    seg2.write_bytes(b"segment 1")

    mock_proc = MagicMock()
    mock_proc.returncode = 1  # Simulated crash

    session = RecordingSession(
        video_id=uuid4(),
        device_id="emulator-5554",
        start_time=time.time() - 30.0,
        data_engine_start_time=time.time() - 30.0,
        local_video_path=seg2,
        android_video_segments=[seg1],
        android_segment_records=[
            {
                "path": seg1,
                "output_path": tmp_path / "recording.mp4",
                "start": 0.0,
                "end": 15.0,
            }
        ],
        android_segment_index=1,
        process=mock_proc,
    )
    set_active_session("emulator-5554", session)

    with (
        patch(
            "artemis.controllers.unified_controller.remux_recording_to_mp4",
            AsyncMock(side_effect=lambda src, dst: (dst.write_bytes(b"final mp4"), True)[1]),
        ),
        patch(
            "artemis.controllers.unified_controller.write_recording_manifest",
            AsyncMock(return_value=tmp_path / "recording.json"),
        ),
    ):
        res = await controller.stop_video_recording()

        assert res.success is True
        assert res.video_path is not None
        assert res.video_path.exists()
        assert str(res.video_path).endswith("recording.mp4")
        assert get_active_session("emulator-5554") is None


@pytest.mark.asyncio
async def test_unified_controller_timeline_alignment(mock_ctx, tmp_path):
    """Verify that relative time offset between DataEngine T0 and video start is precisely calculated."""
    controller = UnifiedMobileController(mock_ctx)
    remove_active_session("emulator-5554")

    mkv_path = tmp_path / "recording.mkv"
    mkv_path.write_bytes(b"dummy mkv content")

    now = time.time()
    # Assume session started 15s ago, and video started 12s ago (offset = 3.0s)
    session_start_t0 = now - 15.0
    video_start_t = now - 12.0

    session = RecordingSession(
        video_id=uuid4(),
        device_id="emulator-5554",
        start_time=video_start_t,
        data_engine_start_time=session_start_t0,
        local_video_path=mkv_path,
        process=MagicMock(returncode=None),
    )
    set_active_session("emulator-5554", session)

    trimmed_ranges = []

    async def mock_trim(src, s_time, e_time, dst):
        trimmed_ranges.append((s_time, e_time))
        dst.write_bytes(b"trimmed")
        return True

    with patch(
        "artemis.controllers.unified_controller.render_timeline_clip",
        AsyncMock(side_effect=mock_trim),
    ):
        # Agent asks for system time range [5.0s, 10.0s]
        res = await controller.extract_segment_metadata(start_time=5.0, end_time=10.0)

        assert res.success is True
        # Since offset = 3.0s, video_start_relative_time should be 5.0 - 3.0 = 2.0s
        # and video_end_relative_time should be 10.0 - 3.0 = 7.0s
        assert len(trimmed_ranges) == 1
        s_time, e_time = trimmed_ranges[0]
        assert abs(s_time - 2.0) < 0.1
        assert abs(e_time - 7.0) < 0.1
        # actual_start_relative_time reported back to agent must align with system T0 (5.0s)
        assert abs(res.actual_start_relative_time - 5.0) < 0.1

    remove_active_session("emulator-5554")