import asyncio
import base64
import io
import json
import stat
import time
import zipfile

import httpx
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from itsdangerous import TimestampSigner
from PIL import Image

import submission_backend as backend


def archive_bytes(name='images/0.png'):
    image = io.BytesIO()
    Image.new('RGB', (8, 8), 'blue').save(image, format='PNG')
    data = io.BytesIO()
    with zipfile.ZipFile(data, 'w') as archive:
        archive.writestr(name, image.getvalue())
        archive.writestr('outputs.jsonl', '{"id": 1}\n')
    return data.getvalue()


def test_model_owner_and_url():
    assert backend.model_id_from_url('https://huggingface.co/alice/model', 'alice') == 'alice/model'
    for value in ['https://huggingface.co/bob/model', 'https://evil.test/alice/model', 'https://huggingface.co@evil.test/alice/model', 'https://huggingface.co/datasets/alice/model', 'https://huggingface.co/alice/model?x=1']:
        with pytest.raises(ValueError):
            backend.model_id_from_url(value, 'alice')


@pytest.mark.parametrize('changes', [{'private': True}, {'gated': 'auto'}, {'cardData': {'license': 'other'}}, {'siblings': []}])
def test_model_rejects_private_gated_nonopen_or_missing_weights(changes):
    info = {'id': 'alice/model', 'private': False, 'gated': False, 'cardData': {'license': 'apache-2.0'}, 'siblings': [{'rfilename': 'model.safetensors'}], **changes}
    async def run():
        async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, json=info))) as client:
            with pytest.raises(ValueError):
                await backend.validate_model('alice/model', client)
    asyncio.run(run())


def test_model_check_is_anonymous():
    def respond(request):
        assert request.headers.get('authorization') == ''
        return httpx.Response(200, json={'id': 'alice/model', 'private': False, 'gated': False, 'sha': 'abc', 'cardData': {'license': 'mit'}, 'siblings': [{'rfilename': 'model.safetensors'}]})
    async def run():
        async with httpx.AsyncClient(transport=httpx.MockTransport(respond), headers={'Authorization': 'Bearer must-not-be-used'}) as client:
            assert (await backend.validate_model('alice/model', client))['revision'] == 'abc'
    asyncio.run(run())


def test_valid_zip(tmp_path):
    path = tmp_path / 'test.zip'
    path.write_bytes(archive_bytes())
    assert backend.inspect_zip(path)['image_count'] == 1


@pytest.mark.parametrize('name', ['../bad.png', '/bad.png', 'C:/bad.png', 'foo\\bad.png', 'run.py'])
def test_unsafe_zip_entries(tmp_path, name):
    path = tmp_path / 'test.zip'
    path.write_bytes(archive_bytes(name))
    with pytest.raises(ValueError):
        backend.inspect_zip(path)


def test_invalid_and_empty_zip(tmp_path):
    path = tmp_path / 'test.zip'
    path.write_text('not really zip')
    with pytest.raises(ValueError):
        backend.inspect_zip(path)
    with zipfile.ZipFile(path, 'w') as archive:
        archive.writestr('metadata.json', '{}')
    with pytest.raises(ValueError, match='none were found'):
        backend.inspect_zip(path)


def test_fake_image_and_symlink(tmp_path):
    path = tmp_path / 'test.zip'
    with zipfile.ZipFile(path, 'w') as archive:
        archive.writestr('fake.png', b'not image')
    with pytest.raises(ValueError):
        backend.inspect_zip(path)
    with zipfile.ZipFile(path, 'w') as archive:
        item = zipfile.ZipInfo('link.png')
        item.create_system = 3
        item.external_attr = (stat.S_IFLNK | 0o777) << 16
        archive.writestr(item, 'target')
    with pytest.raises(ValueError):
        backend.inspect_zip(path)


def cookie(username='alice', csrf='csrf-test'):
    value = {'hf_user': {'username': username, 'sub': 'hf-'+username, 'expires_at': time.time()+3600}, 'csrf': csrf}
    raw = base64.b64encode(json.dumps(value).encode())
    return TimestampSigner('test-secret').sign(raw).decode()


@pytest.fixture
def service(tmp_path, monkeypatch):
    monkeypatch.setenv('SESSION_SECRET', 'test-secret')
    monkeypatch.setenv('SUBMISSIONS_DIR', str(tmp_path / 'private'))
    monkeypatch.setenv('INTREAL_LOCAL_STORAGE', '1')
    monkeypatch.delenv('SPACE_ID', raising=False)
    monkeypatch.delenv('SPACE_HOST', raising=False)
    app = FastAPI()
    backend.install_submissions(app)
    async def verified(repo_id, client):
        return {'repo_id': repo_id, 'revision': 'abc', 'license': 'mit'}
    monkeypatch.setattr(backend, 'validate_model', verified)
    with TestClient(app) as client:
        yield client, tmp_path / 'private'


def form(data):
    return {'model_url': 'https://huggingface.co/alice/model', 'username': 'alice', 'email': 'alice@example.com', 'filename': 'outputs.zip', 'size': len(data)}


def sign_in(client, username='alice'):
    client.cookies.clear()
    client.cookies.set('intreal_session', cookie(username))


def test_auth_and_csrf_required(service):
    client, _ = service
    assert client.post('/api/submissions', json=form(archive_bytes())).status_code == 401
    sign_in(client)
    assert client.post('/api/submissions', json=form(archive_bytes())).status_code == 403
    assert client.put('/api/submissions/'+'a'*32+'/archive', content=b'zip').status_code == 403


def test_forged_username_rejected(service):
    client, _ = service
    sign_in(client)
    payload = form(archive_bytes()) | {'username': 'bob'}
    response = client.post('/api/submissions', json=payload, headers={'X-CSRF-Token':'csrf-test'})
    assert response.status_code == 422


def test_receive_private_archive_and_no_public_download(service):
    client, root = service
    sign_in(client)
    data = archive_bytes()
    headers = {'X-CSRF-Token': 'csrf-test'}
    prepared = client.post('/api/submissions', json=form(data), headers=headers)
    assert prepared.status_code == 200, prepared.text
    sid = prepared.json()['id']
    sign_in(client, 'bob')
    assert client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers).status_code == 404
    sign_in(client)
    received = client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers)
    assert received.status_code == 200, received.text
    assert received.json()['status'] == 'pending_review'
    assert (root / sid / 'outputs.zip').read_bytes() == data
    record = json.loads((root / sid / 'metadata.json').read_text())
    assert record['username'] == 'alice' and record['archive']['image_count'] == 1
    assert not (root / sid / 'images').exists()
    assert client.get(f'/api/submissions/{sid}').status_code == 405
    assert client.get(f'/leaderboard-assets/{sid}/outputs.zip').status_code == 404
    assert client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers).status_code == 409


def test_bad_archive_cleanup(service):
    client, root = service
    sign_in(client)
    data = b'fake zip'
    headers = {'X-CSRF-Token': 'csrf-test'}
    sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
    assert client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers).status_code == 422
    assert not list((root / sid).glob('*.part'))
    assert not (root / sid / 'outputs.zip').exists()
    assert json.loads((root / sid / 'metadata.json').read_text())['status'] == 'failed'


def test_truncated_upload_and_size_limit(service):
    client, root = service
    sign_in(client)
    data = archive_bytes()
    headers = {'X-CSRF-Token': 'csrf-test'}
    too_large = form(data) | {'size': backend.MAX_ZIP + 1}
    assert client.post('/api/submissions', json=too_large, headers=headers).status_code == 422
    sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
    assert client.put(f'/api/submissions/{sid}/archive', content=data[:-10], headers=headers).status_code == 422
    assert not (root / sid / 'outputs.zip').exists()


def test_no_persistence_fails_closed(tmp_path, monkeypatch):
    monkeypatch.setenv('SPACE_ID', 'owner/space')
    monkeypatch.setenv('INTREAL_LOCAL_STORAGE', '1')
    monkeypatch.setenv('SUBMISSIONS_DIR', str(tmp_path))
    monkeypatch.setenv('SUBMISSIONS_MOUNT', str(tmp_path / 'not-a-mount'))
    assert backend.SubmissionService().ready() is False


def test_forged_cookie_is_not_identity(service):
    client, _ = service
    client.cookies.set('intreal_session', cookie()[:-8] + 'forged!!')
    assert client.post('/api/submissions', json=form(archive_bytes()), headers={'X-CSRF-Token': 'csrf-test'}).status_code == 401


def test_duplicate_paths_and_expansion_limit(tmp_path, monkeypatch):
    path = tmp_path / 'test.zip'
    with zipfile.ZipFile(path, 'w') as archive:
        archive.writestr('same.txt', 'a')
        with pytest.warns(UserWarning):
            archive.writestr('same.txt', 'b')
    with pytest.raises(ValueError, match='duplicate'):
        backend.inspect_zip(path)
    path.write_bytes(archive_bytes())
    monkeypatch.setattr(backend, 'MAX_EXPANDED', 1)
    with pytest.raises(ValueError, match='expanded-size'):
        backend.inspect_zip(path)


def test_single_pending_and_weekly_limit(service):
    client, _ = service
    sign_in(client)
    headers = {'X-CSRF-Token': 'csrf-test'}
    data = archive_bytes()
    sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
    assert client.post('/api/submissions', json=form(data), headers=headers).status_code == 409
    assert client.delete('/api/submissions/'+sid, headers=headers).status_code == 200
    for _ in range(3):
        sid = client.post('/api/submissions', json=form(b'invalid'), headers=headers).json()['id']
        assert client.put(f'/api/submissions/{sid}/archive', content=b'invalid', headers=headers).status_code == 422
    assert client.post('/api/submissions', json=form(data), headers=headers).status_code == 429


def test_server_quota_includes_reserved_uploads(service, monkeypatch):
    client, _ = service
    sign_in(client)
    data = archive_bytes()
    monkeypatch.setattr(backend, 'MAX_TOTAL', len(data)-1)
    response = client.post('/api/submissions', json=form(data), headers={'X-CSRF-Token': 'csrf-test'})
    assert response.status_code == 507


def test_invalid_email_and_non_zip_name(service):
    client, _ = service
    sign_in(client)
    for patch in [{'email': 'not-email'}, {'filename': 'outputs.tar'}, {'model_url': 'https://huggingface.co/bob/model'}]:
        response = client.post('/api/submissions', json=form(archive_bytes()) | patch, headers={'X-CSRF-Token':'csrf-test'})
        assert response.status_code == 422


def test_login_disabled_without_real_oauth(service):
    client, _ = service
    assert client.get('/api/submission-session').json()['auth_enabled'] is False
    assert client.get('/auth/login').status_code == 503


def test_full_storage_message_and_disabled_state(service, monkeypatch):
    client, _ = service
    sign_in(client)
    monkeypatch.setattr(backend, 'MAX_TOTAL', 0)
    state = client.get('/api/submission-session').json()
    assert state['storage_full'] is True
    assert 'shaohanyu628@gmail.com' in state['storage_message']
    response = client.post('/api/submissions', json=form(archive_bytes()), headers={'X-CSRF-Token':'csrf-test'})
    assert response.status_code == 507
    assert 'submit again' in response.json()['detail']


def test_low_disk_blocks_upload_before_bytes(service, monkeypatch):
    from collections import namedtuple
    client, _ = service
    sign_in(client)
    disk = namedtuple('usage', 'total used free')
    monkeypatch.setattr(backend.shutil, 'disk_usage', lambda _: disk(10, 9, 1))
    response = client.post('/api/submissions', json=form(archive_bytes()), headers={'X-CSRF-Token':'csrf-test'})
    assert response.status_code == 507
    assert 'shaohanyu628@gmail.com' in response.json()['detail']


def test_disk_fills_during_upload_cleans_partial(service, monkeypatch):
    from collections import namedtuple
    client, root = service
    sign_in(client)
    data = archive_bytes()
    headers = {'X-CSRF-Token':'csrf-test'}
    sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
    disk = namedtuple('usage', 'total used free')
    monkeypatch.setattr(backend.shutil, 'disk_usage', lambda _: disk(10, 9, 1))
    response = client.put(f'/api/submissions/{sid}/archive',content=data,headers=headers)
    assert response.status_code == 507
    assert 'shaohanyu628@gmail.com' in response.json()['detail']
    assert not (root/sid/'outputs.zip.part').exists()
    assert not (root/sid/'outputs.zip').exists()


def test_model_name_is_derived_and_not_user_controlled(service):
    client, root = service
    sign_in(client)
    payload = form(archive_bytes()) | {'model_name': 'forged display name'}
    response = client.post('/api/submissions', json=payload, headers={'X-CSRF-Token':'csrf-test'})
    assert response.status_code == 200
    record = json.loads((root / response.json()['id'] / 'metadata.json').read_text())
    assert record['model_name'] == 'model'


def test_submission_quota_uses_rolling_seven_days(service):
    client, root = service
    sign_in(client)
    headers = {'X-CSRF-Token':'csrf-test'}
    data = archive_bytes()
    for index in range(3):
        sid = f'{index:032x}'
        folder = root / sid
        folder.mkdir(parents=True)
        (folder / 'metadata.json').write_text(json.dumps({
            'id': sid, 'username': 'alice', 'created_at': time.time() - 6*86400,
            'status': 'failed', 'declared_bytes': 10,
        }))
    response = client.post('/api/submissions', json=form(data), headers=headers)
    assert response.status_code == 429
    assert '7-day' in response.json()['detail']
    for path in root.glob('*/metadata.json'):
        record = json.loads(path.read_text())
        record['created_at'] = time.time() - 8*86400
        path.write_text(json.dumps(record))
    assert client.post('/api/submissions', json=form(data), headers=headers).status_code == 200