import urllib.request
import urllib.error
import json
import time

BASE = 'https://dev-crafterx-preference-lab.hf.space'

PASS = 0
FAIL = 0

def post(url, data={}):
    body = json.dumps(data).encode()
    req = urllib.request.Request(url, data=body,
          headers={'Content-Type': 'application/json'})
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            return json.loads(r.read())
    except Exception as e:
        return {'error': str(e)}

def get(url):
    try:
        with urllib.request.urlopen(url, timeout=60) as r:
            return json.loads(r.read())
    except Exception as e:
        return {'error': str(e)}

def get_obs(r):
    """reward/done are at TOP LEVEL, not inside observation."""
    obs = r.get('observation', r)
    obs['reward'] = r.get('reward')
    obs['done']   = r.get('done')
    return obs

def check(name, condition, got=None):
    global PASS, FAIL
    if condition:
        print(f'  ✅ PASS — {name}')
        PASS += 1
    else:
        print(f'  ❌ FAIL — {name} | got: {got}')
        FAIL += 1

print()
print('=' * 60)
print('  PreferenceLab — Full Simulation Test')
print(f'  Target: {BASE}')
print('=' * 60)

# ── TEST 1: Health ─────────────────────────────────────────────
print('\n[1] HEALTH CHECK')
r = get(f'{BASE}/health')
check('Returns status field', 'status' in r, r)
check('Status is healthy', r.get('status') == 'healthy', r.get('status'))

# ── TEST 2: Reset — all 3 task types ──────────────────────────
print('\n[2] RESET — ALL 3 TASK TYPES')
for task in ['pairwise', 'likert', 'consistency']:
    r = post(f'{BASE}/reset', {'task_type': task})
    obs = get_obs(r)
    check(f'reset({task}) returns observation', 'observation' in r, list(r.keys()))
    check(f'reset({task}) has prompt', bool(obs.get('prompt')), obs.get('prompt','')[:30])
    check(f'reset({task}) reward=0.0', obs.get('reward') == 0.0, obs.get('reward'))
    check(f'reset({task}) done=false', obs.get('done') == False, obs.get('done'))
    check(f'reset({task}) task_type correct', obs.get('task_type') == task, obs.get('task_type'))

# ── TEST 3: Pairwise — all choices ────────────────────────────
print('\n[3] PAIRWISE — ALL CHOICES')
rewards_pairwise = {}
for choice in ['A', 'B', 'tie', 'skip']:
    post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': 42})
    r = post(f'{BASE}/step', {'action': {'choice': choice}})
    obs = get_obs(r)
    reward = obs.get('reward')
    rewards_pairwise[choice] = reward
    check(f'choice={choice} reward in [0,1]',
          reward is not None and 0.0 <= reward <= 1.0, reward)

check('A and B give different rewards (anti-constant)',
      rewards_pairwise.get('A') != rewards_pairwise.get('B'),
      f"A={rewards_pairwise.get('A')} B={rewards_pairwise.get('B')}")

check('skip gives partial credit (0.3)',
      rewards_pairwise.get('skip') == 0.3,
      rewards_pairwise.get('skip'))

# ── TEST 4: Likert scoring ─────────────────────────────────────
print('\n[4] LIKERT SCORING')
post(f'{BASE}/reset', {'task_type': 'likert', 'seed': 42})
r = post(f'{BASE}/step', {'action': {
    'helpfulness': 5, 'honesty': 5,
    'harmlessness': 5, 'instruction_following': 5
}})
obs = get_obs(r)
check('Likert perfect scores accepted', obs.get('reward') is not None, list(obs.keys()))
check('Likert reward in [0,1]',
      obs.get('reward') is not None and 0.0 <= obs.get('reward', -1) <= 1.0,
      obs.get('reward'))

post(f'{BASE}/reset', {'task_type': 'likert', 'seed': 42})
r2 = post(f'{BASE}/step', {'action': {
    'helpfulness': 1, 'honesty': 1,
    'harmlessness': 1, 'instruction_following': 1
}})
obs2 = get_obs(r2)
check('Likert worst != perfect (varies)',
      obs.get('reward') != obs2.get('reward'),
      f"perfect={obs.get('reward')} worst={obs2.get('reward')}")

# ── TEST 5: Consistency ranking ────────────────────────────────
print('\n[5] CONSISTENCY RANKING')
post(f'{BASE}/reset', {'task_type': 'consistency', 'seed': 42})
r = post(f'{BASE}/step', {'action': {'ranking': ['A', 'B', 'C', 'D']}})
obs = get_obs(r)
check('Consistency accepts 4-item ranking', obs.get('reward') is not None, list(obs.keys()))
check('Consistency reward in [0,1]',
      obs.get('reward') is not None and 0.0 <= obs.get('reward', -1) <= 1.0,
      obs.get('reward'))

post(f'{BASE}/reset', {'task_type': 'consistency', 'seed': 42})
r2 = post(f'{BASE}/step', {'action': {'ranking': ['D', 'C', 'B', 'A']}})
obs2 = get_obs(r2)
check('Consistency reversed != perfect (grader works)',
      obs.get('reward') != obs2.get('reward'),
      f"correct={obs.get('reward')} reversed={obs2.get('reward')}")

# ── TEST 6: Full episode — 5 steps ────────────────────────────
print('\n[6] FULL EPISODE — 5 STEPS')
post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': 99})
rewards = []
done = False
steps = 0
for i in range(10):
    r = post(f'{BASE}/step', {'action': {'choice': 'A'}})
    obs = get_obs(r)
    rewards.append(obs.get('reward', 0))
    done = obs.get('done', False)
    steps += 1
    if done:
        break

check('Episode terminates (done=true)', done == True, done)
check('Episode runs exactly 5 steps', steps == 5, steps)
check('All rewards in [0,1]',
      all(0.0 <= rv <= 1.0 for rv in rewards if rv is not None), rewards)

# ── TEST 7: Reproducibility with seed ─────────────────────────
print('\n[7] REPRODUCIBILITY — SAME SEED = SAME EPISODE')
r1 = post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': 777})
r2 = post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': 777})
p1 = get_obs(r1).get('prompt', 'X')
p2 = get_obs(r2).get('prompt', 'Y')
check('Same seed produces same prompt', p1 == p2, f"p1={p1[:30]} p2={p2[:30]}")

r3 = post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': 888})
p3 = get_obs(r3).get('prompt', 'Z')
check('Different seed produces different episode',
      p1 != p3, f"seed777={p1[:30]} seed888={p3[:30]}")

# ── TEST 8: State endpoint ─────────────────────────────────────
print('\n[8] STATE ENDPOINT')
post(f'{BASE}/reset', {'task_type': 'likert', 'seed': 42})
post(f'{BASE}/step', {'action': {
    'helpfulness': 4, 'honesty': 4,
    'harmlessness': 4, 'instruction_following': 4
}})
r = get(f'{BASE}/state')
check('State has step_count > 0', r.get('step_count', 0) > 0, r.get('step_count'))
check('State has task_type=likert', r.get('task_type') == 'likert', r.get('task_type'))
check('State has seed=42', r.get('seed') == 42, r.get('seed'))
check('State has cumulative_reward', 'cumulative_reward' in r, list(r.keys()))

# ── TEST 9: Concurrent resets ──────────────────────────────────
print('\n[9] CONCURRENT RESETS — ISOLATION CHECK')
ra = post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': 1})
rb = post(f'{BASE}/reset', {'task_type': 'consistency', 'seed': 2})
ta = get_obs(ra).get('task_type')
tb = get_obs(rb).get('task_type')
check('Different task types accepted',
      ta in ['pairwise','likert','consistency'] and
      tb in ['pairwise','likert','consistency'], f"ta={ta} tb={tb}")

# ── TEST 10: Disqualification guard ───────────────────────────
print('\n[10] DISQUALIFICATION GUARD — REWARDS MUST VARY')
all_rewards = set()
for seed in [1, 2, 3, 4, 5]:
    post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': seed})
    for choice in ['A', 'B']:
        r = post(f'{BASE}/step', {'action': {'choice': choice}})
        obs = get_obs(r)
        rw = obs.get('reward')
        if rw is not None:
            all_rewards.add(rw)
    post(f'{BASE}/reset', {'task_type': 'pairwise', 'seed': seed})

check('Rewards NOT constant (DQ check)',
      len(all_rewards) > 1, f"unique rewards: {all_rewards}")

# ── FINAL SUMMARY ──────────────────────────────────────────────
print()
print('=' * 60)
total = PASS + FAIL
print(f'  RESULTS: {PASS}/{total} PASSED')
print()
if FAIL == 0:
    print('  🏆 ALL TESTS PASSED — SUBMISSION READY!')
elif FAIL <= 3:
    print('  ⚠️  MINOR ISSUES — Check above')
else:
    print('  ❌  FAILURES — Fix before submitting')
print('=' * 60)