import threading
import time

import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient

from server import GenerationRequest, InferenceEngine, create_app, normalize_input


def test_protein_validation_preserves_one_sequence_and_reports_normalization():
    request = GenerationRequest(direction='p2t', input='>candidate\nmaK uZob\nACD')
    assert request.greedy is False
    assert request.temperature == 1
    assert request.max_length == 512
    normalized, note = normalize_input(request)
    assert normalized == 'MAKXXXXACD'
    assert note
    for sequence in ['>one\nMAK\n>two\nMAK', 'MAK123', 'A' * 512, '>header']:
        with pytest.raises(HTTPException) as error:
            normalize_input(GenerationRequest(direction='p2t', input=sequence))
        assert error.value.status_code == 422


def test_api_rejects_invalid_requests_and_does_not_fake_preview_results():
    with TestClient(create_app(InferenceEngine(disabled=True))) as client:
        assert client.get('/api/health').json()['inference'] == 'disabled'
        assert client.post('/api/jobs', json={'direction': 't2p', 'input': 'PET hydrolase'}).status_code == 503
        for payload in [
            {'direction': 'invalid', 'input': 'abc'},
            {'direction': 't2p', 'input': '   '},
            {'direction': 't2p', 'input': 'x', 'max_length': 10000},
            {'direction': 'p2t', 'input': 'MAK9'},
        ]:
            assert client.post('/api/jobs', json=payload).status_code == 422
        assert client.get('/api/jobs/nonexistent').status_code == 404


def test_queue_is_bounded():
    engine = InferenceEngine(disabled=False)
    for _ in range(5):
        engine.submit(GenerationRequest(direction='t2p', input='Hydrolase'))
    with pytest.raises(HTTPException) as error:
        engine.submit(GenerationRequest(direction='t2p', input='Hydrolase'))
    assert error.value.status_code == 429
    assert len(engine.jobs) == 5


class RecordingRunner:
    def __init__(self):
        self.loads = 0
        self.active = 0
        self.max_active = 0
        self.calls = []

    def load(self):
        self.loads += 1

    def generate(self, request, sequence):
        self.active += 1
        self.max_active = max(self.active, self.max_active)
        self.calls.append((request.direction, sequence))
        time.sleep(.03)
        self.active -= 1
        if sequence == 'fail':
            raise ValueError('Please shorten the description.')
        return 'MAK' if request.direction == 't2p' else 'A test annotation.'


def test_jobs_run_serially_reuse_model_and_preserve_input_provenance():
    runner = RecordingRunner()
    with TestClient(create_app(InferenceEngine(runner, disabled=False))) as client:
        first = client.post('/api/jobs', json={'direction': 't2p', 'input': 'Hydrolase'}).json()['id']
        second = client.post('/api/jobs', json={'direction': 'p2t', 'input': '>candidate\nma u'}).json()['id']
        deadline = time.monotonic() + 3
        while time.monotonic() < deadline:
            b = client.get('/api/jobs/' + second).json()
            if b['status'] == 'completed':
                break
            time.sleep(.02)
        a = client.get('/api/jobs/' + first).json()
        assert a['status'] == b['status'] == 'completed'
        assert a['output'] == 'MAK'
        assert b['normalized_input'] == 'MAX'
        assert b['normalization_note']
        assert runner.loads == 1 and runner.max_active == 1
        assert runner.calls == [('t2p', 'Hydrolase'), ('p2t', 'MAX')]


def test_failed_generation_does_not_become_a_success_or_poison_the_model():
    runner = RecordingRunner()
    engine = InferenceEngine(runner, disabled=False)
    engine.start()
    try:
        first = engine.submit(GenerationRequest(direction='t2p', input='fail'))['id']
        second = engine.submit(GenerationRequest(direction='t2p', input='hydrolase'))['id']
        deadline = time.monotonic() + 3
        while time.monotonic() < deadline and engine.read(second)['status'] != 'completed':
            time.sleep(.02)
        assert engine.read(first)['status'] == 'failed'
        assert engine.read(first)['output'] is None
        assert engine.read(second)['status'] == 'completed'
        assert runner.loads == 1
    finally:
        engine.close()