'''
Copyright (C) 2024-2026 Mo Zhou <lumin@debian.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
'''
import socket
from unittest.mock import MagicMock, patch
import pytest
import zmq
from debgpt.backend import ZMQBackend
from debgpt.frontend import ZMQFrontend
from debgpt.defaults import Config
from types import SimpleNamespace
import threading
import time

def find_free_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('', 0))
        return s.getsockname()[1]

class DummyLLM:
    def __init__(self):
        class MockTok:
            def apply_chat_template(self, messages, tokenize=True, return_tensors='pt'):
                return [1, 2, 3]
        self.tok = MockTok()
        self.last_kwargs = {}

    def generate(self, messages, **kwargs):
        self.last_kwargs = kwargs
        import copy
        new_messages = copy.deepcopy(messages)
        new_messages.append({'role': 'assistant', 'content': 'Hello, non-streaming world!'})
        return new_messages

    def generate_stream(self, messages, **kwargs):
        self.last_kwargs = kwargs
        yield "Hello"
        yield ", "
        yield "streaming"
        yield " "
        yield "world!"
    
    def __call__(self, messages, **kwargs):
        return self.generate(messages, **kwargs)

def test_zmq_streaming():
    port = find_free_port()
    host = "tcp://127.0.0.1"
    
    backend_args = SimpleNamespace(
        host="tcp://*",
        port=port,
        backend_impl="zmq",
        model_id="Qwen/Qwen3.5-0.8B",
        device="cpu",
        precision="fp16",
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9
    )

    frontend_args = SimpleNamespace(
        zmq_backend=f"{host}:{port}",
        frontend="zmq",
        debgpt_home="/tmp",
        color='no',
        multiline=False,
        render='no',
        vertical_overflow="visible",
        verbose=False,
        temperature=0.8,
        top_p=0.95
    )

    dummy_llm = DummyLLM()
    with patch('debgpt.llm.create_llm', return_value=dummy_llm):
        backend = ZMQBackend(backend_args)

    def run_backend():
        try:
            backend.server()
        except Exception:
            pass

    backend_thread = threading.Thread(target=run_backend)
    backend_thread.daemon = True
    backend_thread.start()

    time.sleep(0.5)

    try:
        # Test streaming mode
        frontend = ZMQFrontend(frontend_args)
        
        user_message = "Test streaming message"
        response = frontend.query(user_message)
        assert response == "Hello, streaming world!"
        assert len(frontend.session) == 2
        assert frontend.session[0]['role'] == 'user'
        assert frontend.session[0]['content'] == user_message
        assert frontend.session[1]['role'] == 'assistant'
        assert frontend.session[1]['content'] == "Hello, streaming world!"
        
        # Verify that temperature and top_p parameters are correctly forwarded
        assert dummy_llm.last_kwargs.get('temperature') == 0.8
        assert dummy_llm.last_kwargs.get('top_p') == 0.95

    finally:
        pass
