import json
import socket
import unittest
import urllib.error

from tests.test_gitea import FakeOpener, FakeResponse
from werkrevisor.jev import JevClient, JevError, input_cost_usd

KEY = "typesafe-secret-key-0123456789"
ENDPOINT = "https://api.typesafe.ai/v1/systemone"
QUESTIONS = {"vague_then": {"type": "noul", "instructions": "Is it vague?"}}
OK = {"model": "jev-1.13.0", "answers": {"vague_then": {"type": "noul", "noul": 0.2}},
      "usage": {"input_tokens": 306, "output_tokens": 23}}


def http_error(code):
    def raise_it(request):
        raise urllib.error.HTTPError(request.full_url, code, "error", {}, None)
    return raise_it


def ok(request):
    return FakeResponse(json.dumps(OK).encode())


def client(*handlers):
    queue = list(handlers)
    opener = FakeOpener(lambda request: queue.pop(0)(request))
    sleeps = []
    return JevClient(ENDPOINT, KEY, "jev-latest", 7, opener=opener, sleep=sleeps.append), opener, sleeps


class RequestTest(unittest.TestCase):
    def test_request_shape(self):
        c, opener, _ = client(ok)
        answer = c.ask({"scenario": {"id": "1.01", "text": "Ä"}}, QUESTIONS)
        request = opener.requests[0]
        self.assertEqual((request.full_url, request.get_method()), (ENDPOINT, "POST"))
        self.assertEqual(request.get_header("Authorization"), f"Bearer {KEY}")
        self.assertEqual(request.get_header("Content-type"), "application/json")
        self.assertEqual(json.loads(request.data.decode("utf-8")),
                         {"model": "jev-latest", "state": {"scenario": {"id": "1.01", "text": "Ä"}},
                          "questions": QUESTIONS})
        self.assertEqual((answer.model, answer.input_tokens, answer.output_tokens, answer.attempts),
                         ("jev-1.13.0", 306, 23, 1))
        self.assertEqual(answer.answers, OK["answers"])

    def test_cost_counts_input_tokens_only(self):
        self.assertAlmostEqual(input_cost_usd(1_000_000, 0.042), 0.042)
        self.assertEqual(input_cost_usd(0, 0.042), 0)


class RetryTest(unittest.TestCase):
    def test_one_retry_on_5xx(self):
        c, opener, sleeps = client(http_error(529), ok)
        self.assertEqual(c.ask({}, QUESTIONS).attempts, 2)
        self.assertEqual((len(opener.requests), len(sleeps)), (2, 1))

    def test_one_retry_on_timeout(self):
        def timeout(request):
            raise socket.timeout("timed out")
        c, opener, _ = client(timeout, ok)
        self.assertEqual(c.ask({}, QUESTIONS).attempts, 2)

    def test_second_failure_is_an_error(self):
        c, opener, _ = client(http_error(503), http_error(502))
        with self.assertRaisesRegex(JevError, "HTTP 502"):
            c.ask({}, QUESTIONS)
        self.assertEqual(len(opener.requests), 2)

    def test_no_retry_on_client_errors(self):
        for code in (401, 422, 429):
            with self.subTest(code=code):
                c, opener, _ = client(http_error(code))
                with self.assertRaises(JevError) as caught:
                    c.ask({}, QUESTIONS)
                self.assertEqual(len(opener.requests), 1)
                self.assertIn(f"HTTP {code}", str(caught.exception))
                self.assertNotIn(KEY, str(caught.exception))

    def test_connection_error_names_no_key(self):
        def refused(request):
            raise urllib.error.URLError(ConnectionRefusedError("refused"))
        c, opener, _ = client(refused)
        with self.assertRaises(JevError) as caught:
            c.ask({}, QUESTIONS)
        self.assertNotIn(KEY, str(caught.exception))
        self.assertEqual(len(opener.requests), 1)


class AnswerValidationTest(unittest.TestCase):
    def test_invalid_answers(self):
        for payload in (b"<html>", b"[]", json.dumps({"answers": {}}).encode(),
                        json.dumps({"answers": OK["answers"]}).encode(),
                        json.dumps(dict(OK, usage={"input_tokens": -1, "output_tokens": 0})).encode()):
            with self.subTest(payload=payload):
                c, _, _ = client(lambda request, p=payload: FakeResponse(p))
                with self.assertRaises(JevError):
                    c.ask({}, QUESTIONS)


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