"""Exercise the real PowerShell setup downloader without downloading 2.5 GB."""
from pathlib import Path
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import hashlib
import json
import subprocess
import tempfile
import threading

ROOT = Path(__file__).resolve().parents[2]
SCRIPT = ROOT / 'tools/fortune_ai/download_model.ps1'
PAYLOAD = b'GGUF' + bytes(range(256)) * 4096


class Server(BaseHTTPRequestHandler):
    requests = []

    def log_message(self, *_):
        pass

    def do_GET(self):
        self.requests.append((self.path, self.headers.get('Range')))
        if self.path == '/fail':
            self.send_error(503)
            return
        offset = 0
        if self.path == '/resume' and self.headers.get('Range'):
            offset = int(self.headers['Range'].split('=')[1].split('-')[0])
        data = (b'x' * len(PAYLOAD) if self.path == '/corrupt' else PAYLOAD)[offset:]
        self.send_response(206 if offset else 200)
        if offset:
            self.send_header('Content-Range', f'bytes {offset}-{len(PAYLOAD)-1}/{len(PAYLOAD)}')
        self.send_header('Content-Length', str(len(data)))
        self.end_headers()
        self.wfile.write(data)


def main():
    assert SCRIPT.is_file(), 'Missing one-click PowerShell model downloader'
    server = ThreadingHTTPServer(('127.0.0.1', 0), Server)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    try:
        (ROOT / 'build/scratch').mkdir(parents=True, exist_ok=True)
        with tempfile.TemporaryDirectory(prefix='fortune-download-', dir=ROOT / 'build/scratch') as tmp:
            base = Path(tmp)
            for name, routes in [('fallback', ['fail', 'resume']), ('resume', ['resume']),
                                 ('no-range', ['full']), ('checksum', ['corrupt', 'resume']),
                                 ('pointer', ['resume']), ('completed-part', ['fail']),
                                 ('failure', ['fail'])]:
                out = base / name
                out.mkdir()
                manifest = out / 'model.json'
                manifest.write_text(json.dumps({'file': 'models/test.gguf', 'bytes': len(PAYLOAD),
                    'sha256': hashlib.sha256(PAYLOAD).hexdigest(),
                    'download_urls': [f'http://127.0.0.1:{server.server_port}/{r}' for r in routes]}))
                model = out / 'models/test.gguf'
                model.parent.mkdir()
                if name in ('resume', 'no-range'):
                    model.with_suffix('.gguf.part').write_bytes(PAYLOAD[:65536])
                if name == 'failure':
                    model.write_bytes(b'existing invalid file must survive failure')
                if name == 'pointer':
                    model.write_bytes(b'version https://git-lfs.github.com/spec/v1\n')
                if name == 'completed-part':
                    model.with_suffix('.gguf.part').write_bytes(PAYLOAD)
                result = subprocess.run(['powershell.exe', '-NoProfile', '-ExecutionPolicy', 'Bypass',
                    '-File', str(SCRIPT), '-ManifestPath', str(manifest), '-AiDirectory', str(out)],
                    capture_output=True, text=True, timeout=60)
                if name == 'failure':
                    assert result.returncode != 0
                    assert model.read_bytes() == b'existing invalid file must survive failure'
                    continue
                assert result.returncode == 0, result.stdout + result.stderr
                assert model.read_bytes() == PAYLOAD, name
                before = len(Server.requests)
                result = subprocess.run(['powershell.exe', '-NoProfile', '-ExecutionPolicy', 'Bypass',
                    '-File', str(SCRIPT), '-ManifestPath', str(manifest), '-AiDirectory', str(out)],
                    capture_output=True, text=True, timeout=30)
                assert result.returncode == 0 and len(Server.requests) == before, 'Valid model redownloaded'
            assert any(route == '/resume' and rng == 'bytes=65536-' for route, rng in Server.requests)
            assert any(route == '/full' and rng == 'bytes=65536-' for route, rng in Server.requests)
        print('PASS: actual PowerShell download, fallback, resume, ignored Range, SHA-256, safe failure, no-op')
    finally:
        server.shutdown()
        server.server_close()


if __name__ == '__main__':
    main()