import unittest
import requests
import os

class QwenConnectionTests(unittest.TestCase):

    def setUp(self):
        """Setup method to retrieve the Ollama endpoint from environment variables."""
        self.ollama_endpoint = os.environ.get("OLLAMA_ENDPOINT")
        self.assertIsNotNone(self.ollama_endpoint, "OLLAMA_ENDPOINT environment variable must be set.")
        self.model_name = "qwen" # Assuming Qwen is the default model
        self.test_prompt = "What is the capital of France?"

    def test_connection_success(self):
        """Test a successful connection to the Ollama endpoint."""
        try:
            response = requests.get(self.ollama_endpoint, timeout=5)  # Simple GET to check connectivity
            self.assertEqual(response.status_code, 200, f"Connection failed: Status code {response.status_code}")
        except requests.ConnectionError as e:
            self.fail(f"Connection failed: {e}")

    def test_model_availability(self):
        """Test if the specified model is available on the Ollama endpoint."""
        try:
            response = requests.post(f"{self.ollama_endpoint}/api/generate", json={
                "prompt": self.test_prompt,
                "model": self.model_name,
                "stream": False  #  Stream: false for simple non-stream test.
            }, timeout=10)
            self.assertEqual(response.status_code, 200, f"Model unavailable: Status code {response.status_code}.  Response text: {response.text}")

            response_json = response.json()
            self.assertIn("response", response_json, "Response JSON missing 'response' key.")
            self.assertIsInstance(response_json["response"], str, "'response' value should be a string.")
            self.assertGreater(len(response_json["response"]), 0, "'response' should not be empty.")
        except requests.RequestException as e:
            self.fail(f"Request failed: {e}")
        except (KeyError, ValueError) as e:
            self.fail(f"Error parsing response: {e}.  Raw response: {response.text if 'response' in locals() else 'No response'}")

    def test_streaming_connection(self):
        """Test a successful streaming connection to the Ollama endpoint."""
        try:
            response = requests.post(f"{self.ollama_endpoint}/api/generate", json={
                "prompt": self.test_prompt,
                "model": self.model_name,
                "stream": True  #  Stream: True for streaming test.
            }, stream=True, timeout=10)
            self.assertEqual(response.status_code, 200, f"Streaming failed: Status code {response.status_code}")

            # Check the first few lines to verify the stream is working.  Avoid reading the entire stream.
            lines_processed = 0
            for line in response.iter_lines():
                if line:
                    lines_processed += 1
                    # Check if the line is valid JSON and contains a 'response' key.
                    try:
                        json_data = json.loads(line.decode('utf-8'))
                        self.assertIn("response", json_data, "Streaming response missing 'response' key.")
                    except json.JSONDecodeError:
                        self.fail(f"Streaming response is not valid JSON: {line.decode('utf-8')}")
                    except AssertionError as e:
                        self.fail(f"Streaming assertion failed: {e}")

                if lines_processed >= 3:
                    break  # Only check the first few lines

            if lines_processed == 0:
                self.fail("No data received from streaming endpoint.")


        except requests.RequestException as e:
            self.fail(f"Streaming request failed: {e}")
        except Exception as e:
            self.fail(f"Error processing streaming response: {e}")


if __name__ == '__main__':
    import json
    unittest.main()