import pytest
import uuid
import asyncio
from channels.testing import WebsocketCommunicator
from project.config.asgi import application
from project.remotectl.models import RemoteHost, BatchScript, CommandTask

class DummyConn:
    def close(self):
        pass

@pytest.mark.django_db(transaction=True)
@pytest.mark.asyncio
async def test_command_failure(monkeypatch, django_user_model):
    user = await django_user_model.objects.acreate(username='userf')
    host = await RemoteHost.objects.acreate(name='hf', hostname='localhost', username='u')

    from project.remotectl import consumers

    async def fake_open_connection(rh):
        return DummyConn()

    async def fake_run_command(conn, command, on_chunk, cancel_event):
        await on_chunk('stdout', 'doing something\n')
        return 5  # non-zero

    monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
    monkeypatch.setattr(consumers, 'run_command', fake_run_command)

    session_id = uuid.uuid4()
    comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
    comm.scope['user'] = user
    connected, _ = await comm.connect(); assert connected
    await comm.receive_json_from()  # connected
    await comm.send_json_to({'action':'start','host_id':host.id,'command':'failingcmd'})
    msg = await comm.receive_json_from(); assert msg['event']=='started'
    # drain until completed
    status = None
    while True:
        m = await comm.receive_json_from()
        if m['event'] == 'completed':
            status = m['status']; break
    assert status == 'failed'
    await comm.disconnect()

@pytest.mark.django_db(transaction=True)
@pytest.mark.asyncio
async def test_batch_success(monkeypatch, django_user_model):
    user = await django_user_model.objects.acreate(username='userb')
    host = await RemoteHost.objects.acreate(name='hb', hostname='localhost', username='u')
    batch = await BatchScript.objects.acreate(name='b1', description='d', script='cd /tmp\necho one\necho two\n# comment\necho three')

    from project.remotectl import consumers
    async def fake_open_connection(rh):
        return DummyConn()

    calls = []
    async def fake_run_command(conn, command, on_chunk, cancel_event):
        calls.append(command)
        await on_chunk('stdout', 'out\n')
        return 0

    monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
    monkeypatch.setattr(consumers, 'run_command', fake_run_command)

    session_id = uuid.uuid4()
    comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
    comm.scope['user'] = user
    connected, _ = await comm.connect(); assert connected
    await comm.receive_json_from()  # connected
    await comm.send_json_to({'action':'start_batch','host_id':host.id,'batch_id':batch.id})
    msg = await comm.receive_json_from(); assert msg['event']=='started'
    step_headers = 0
    completed = None
    while True:
        m = await comm.receive_json_from()
        if m['event']=='chunk' and m['data'].startswith('\n>>>'):
            step_headers += 1
        if m['event']=='completed':
            completed = m; break
    assert completed['status']=='ok'
    # Steps: cd, echo one, echo two, echo three => 4 headers
    assert step_headers == 5 or step_headers >= 4  # tolerate extra formatting
    # run_command should have been called for 3 executable echo commands
    assert len(calls) == 3
    await comm.disconnect()

@pytest.mark.django_db(transaction=True)
@pytest.mark.asyncio
async def test_batch_failure(monkeypatch, django_user_model):
    user = await django_user_model.objects.acreate(username='userbf')
    host = await RemoteHost.objects.acreate(name='hbf', hostname='localhost', username='u')
    batch = await BatchScript.objects.acreate(name='bf', description='d', script='echo ok\necho fail\necho skip')

    from project.remotectl import consumers
    async def fake_open_connection(rh):
        return DummyConn()

    call_count = 0
    async def fake_run_command(conn, command, on_chunk, cancel_event):
        nonlocal call_count
        call_count += 1
        await on_chunk('stdout', f'run {call_count}\n')
        if call_count == 2:
            return 9  # fail second step
        return 0

    monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
    monkeypatch.setattr(consumers, 'run_command', fake_run_command)

    session_id = uuid.uuid4()
    comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
    comm.scope['user'] = user
    connected, _ = await comm.connect(); assert connected
    await comm.receive_json_from()  # connected
    await comm.send_json_to({'action':'start_batch','host_id':host.id,'batch_id':batch.id})
    await comm.receive_json_from()  # started
    while True:
        m = await comm.receive_json_from()
        if m['event']=='completed':
            assert m['status'] == 'failed'
            break
    assert call_count == 2  # third not executed
    await comm.disconnect()

@pytest.mark.django_db(transaction=True)
@pytest.mark.asyncio
async def test_cancel_command(monkeypatch, django_user_model):
    user = await django_user_model.objects.acreate(username='userc')
    host = await RemoteHost.objects.acreate(name='hc', hostname='localhost', username='u')

    from project.remotectl import consumers
    async def fake_open_connection(rh):
        return DummyConn()

    async def fake_run_command(conn, command, on_chunk, cancel_event):
        # emit some output repeatedly until canceled
        for i in range(5):
            await on_chunk('stdout', f'line {i}\n')
            await asyncio.sleep(0)
            if cancel_event.is_set():
                return 130
        return 0

    monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
    monkeypatch.setattr(consumers, 'run_command', fake_run_command)

    session_id = uuid.uuid4()
    comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
    comm.scope['user'] = user
    connected, _ = await comm.connect(); assert connected
    await comm.receive_json_from()  # connected
    await comm.send_json_to({'action':'start','host_id':host.id,'command':'longrun'})
    await comm.receive_json_from()  # started
    # send cancel quickly
    await comm.send_json_to({'action':'cancel'})
    saw_canceling = False
    status = None
    while True:
        m = await comm.receive_json_from()
        if m['event']=='canceling':
            saw_canceling = True
        if m['event']=='completed':
            status = m['status']; break
    assert saw_canceling
    assert status == 'canceled'
    await comm.disconnect()