import time
import pytest
from antler.generators import Generator
from antler.harnesses import Harness
from antler.harnesses.linearharness import LinearHarness
from antler.harnesses.multiprocessharness import MultiProcessHarness
from antler.explorers.exhaustivesearch import ExhaustiveSearch
from antler.techniques.acceptingprefix import AcceptingPrefix
from antler.techniques.addnoise import AddNoise
from antler.techniques.encoding import Encoding
from antler.techniques.refusalsuppression import RefusalSuppression
from antler.probes.cursewordfuck import CurseWordFuck
from antler.probes.illegaldrugs import IllegalDrugs
from antler.attempt import Attempt
from antler.transforms.transform import Transform

class SleepingGenerator(Generator):
    def generate(self, prompt: str) -> str:
        time.sleep(0.5) # Simulate expensive task, while possibly allowing other threads to still work
        return "Heyo this is the best answer you could hope for"
    
class SometimesFailingGenerator(Generator):
    def __init__(self, model: str, key: str = None, options: dict = ...) -> None:
        super().__init__(model, key, options)
        self.fail = True

    def generate(self, prompt: str) -> str:
        if self.fail:
            self.fail = False
            raise(TimeoutError("The connection timed out!"))
        else:
            self.fail = True
            return "This is an answer that might work"
        
class Base64OfSpacesGenerator(Generator):
    def generate(self, prompt: str) -> str:
        return "ZiB1IGMgaw==" #base64 of the string "f u c k"
        

def run_linear_harness():
    harness = LinearHarness([CurseWordFuck(), IllegalDrugs()], ExhaustiveSearch([AcceptingPrefix(), RefusalSuppression()]), SleepingGenerator, "key", "123")
    harness.run()

def run_multiprocess_harness():
    harness = MultiProcessHarness([CurseWordFuck(), IllegalDrugs()], ExhaustiveSearch([AcceptingPrefix(), RefusalSuppression()]), SleepingGenerator, "key", "123")
    harness.run()

@pytest.fixture
def linear_harness():
    return LinearHarness([CurseWordFuck(), IllegalDrugs()], ExhaustiveSearch([AcceptingPrefix(), RefusalSuppression()]), SleepingGenerator, "key", "123")

def test_all_attempts_correct_type(linear_harness):
    linear_harness.log_writer.LogRunParams({
        "generator_type": "Testing type",
        "model": "Test model",
        "model_options": {},
        "repetitions": 1,
    })
    attempts = linear_harness.collectAttempts()
    for attempt in attempts:
        assert type(attempt) == Attempt


def test_collect_not_implemented():
    with pytest.raises(NotImplementedError):
        harness = Harness([], None, SleepingGenerator, "key", "model")
        harness.collectAttempts()

def test_linear_time_for_linear_harness():
    start = time.time()
    run_linear_harness()
    end = time.time()
    process_time = end - start # Number of seconds the whole harness run took in total
    assert(process_time > 4) # There are 4 transforms generated by the exhaustive search explorer with 2 techniques * 2 probes at 0.5 seconds each
    assert(process_time < 6) # Some wiggle room

def test_faster_multiprocess_time():
    start = time.time()
    run_multiprocess_harness()
    end = time.time()
    process_time_1 = end - start
    start = time.time()
    run_linear_harness()
    end = time.time()
    process_time_2 = end - start
    assert(process_time_1 < process_time_2)

@pytest.mark.parametrize("attempts, result", [
    (
        [
            Attempt(Transform([ AcceptingPrefix() ]), CurseWordFuck()),
            Attempt(Transform([ AcceptingPrefix() ]), CurseWordFuck())
        ],
        [
            Attempt(Transform([ AcceptingPrefix() ]), CurseWordFuck())
        ]
    )
])
def test_collapse_attempts(attempts: list[Attempt], result: list[Attempt]):
    attempts[0].addResponseObject({"answer": "fuck", "detection": True})
    attempts[1].addResponseObject({"answer": "no", "detection": False})
    result[0].addResponseObject({"answer": "fuck", "detection": True})
    result[0].addResponseObject({"answer": "no", "detection": False})
    collapsed_attempts = Harness.collapseSameAttempts(attempts)
    assert len(collapsed_attempts) == len(result)
    for i in range(len(collapsed_attempts)):
        assert collapsed_attempts[i] == result[i]

def test_error_handling_harness():
    harness = LinearHarness([CurseWordFuck(), IllegalDrugs()], ExhaustiveSearch([AcceptingPrefix()]), SometimesFailingGenerator, "key", "123", repetitions=3)
    harness.run()

def test_harness_answer_fixing():
    harness = LinearHarness([], ExhaustiveSearch([]), Base64OfSpacesGenerator, "key", "123")
    attempt: Attempt = Harness.runAttempt((harness, Attempt(Transform([AddNoise(), Encoding()]), CurseWordFuck())))
    print(attempt.getReplies())
    assert attempt.getAttemptSuccessRate() == 1