import io,json,os,threading,time,uuid
from pathlib import Path
from datetime import datetime,timezone
import gradio as gr
from huggingface_hub import HfApi
ROOT=Path(__file__).parent
QS=json.loads((ROOT/'questions.json').read_text())
LOCK=threading.Lock()
CRITERIA='''## Judging criteria
**At least one option fits in every question.** Two or all three may also fit. For each image, choose **Fits** or **Does not fit**.

- **Surface fit:** the object’s entire footprint must stay within the usable top surface.
- **Container fit:** the entire rigid object must fit inside the container without protruding.
'''
def new_state():return {'id':str(uuid.uuid4()),'index':0,'answers':{},'start':time.time(),'submitted':False}
def render(s):
 q=QS[s['index']];a=s['answers'].get(q['id'],{})
 return [s,f"### Question {s['index']+1} of 24 · {'Surface fit' if q['domain']=='area' else 'Container fit'}\n\n{q['question']}",*[str(ROOT/'images'/Path(q['images'][l]).name) for l in 'ABC'],*[gr.update(value=a.get(l),interactive=not s['submitted']) for l in 'ABC'],gr.update(interactive=s['index']>0 and not s['submitted']),gr.update(value='Submit' if s['index']==23 else 'Next',interactive=not s['submitted']), '']
def init():return render(new_state())
def move(s,a,b,c,direction):
 if s is None:s=new_state()
 if s['submitted']:raise gr.Error('Your responses have already been submitted.')
 if direction==1 and any(v not in ['Fits','Does not fit'] for v in [a,b,c]):raise gr.Error('Please judge A, B, and C before continuing.')
 s['answers'][QS[s['index']]['id']]=dict(zip('ABC',[a,b,c]))
 if direction==1 and s['index']==23:
  if len(s['answers'])!=24 or any(any(v not in ['Fits','Does not fit'] for v in r.values()) for r in s['answers'].values()):raise gr.Error('Please answer all options before submitting.')
  token,repo=os.getenv('HF_TOKEN'),os.getenv('RESPONSES_REPO')
  if not token or not repo:raise gr.Error('Response storage is not configured. Please contact the study owner.')
  data={'participant_id':s['id'],'study_version':'fit-min1-24-20261009-v1','submitted_at':datetime.now(timezone.utc).isoformat(),'elapsed_seconds':int(time.time()-s['start']),'answers':{qid:{l:'Yes' if v=='Fits' else 'No' for l,v in row.items()} for qid,row in s['answers'].items()}}
  try:
   with LOCK:HfApi(token=token).upload_file(path_or_fileobj=io.BytesIO(json.dumps(data,indent=2).encode()),path_in_repo=f"responses/{s['id']}.json",repo_id=repo,repo_type='dataset',commit_message='Save FitTest submission')
  except Exception:raise gr.Error('Your submission was not saved. Please press Submit again.')
  s['submitted']=True;out=render(s);out[-1]=f"## Thank you!\nYour responses have been saved privately.\n\nReceipt: `{s['id']}`";return out
 s['index']=max(0,min(23,s['index']+direction));return render(s)
def previous(s,a,b,c):return move(s,a,b,c,-1)
def next_question(s,a,b,c):return move(s,a,b,c,1)
with gr.Blocks(title='FitTest') as demo:
 gr.Markdown('# FitTest')
 gr.Markdown(CRITERIA)
 state=gr.State()
 question=gr.Markdown()
 imgs=[];radios=[]
 with gr.Row():
  for l in 'ABC':
   with gr.Column():
    imgs.append(gr.Image(label='Option '+l,interactive=False))
    radios.append(gr.Radio(['Fits','Does not fit'],label=l,value=None))
 with gr.Row():
  prev=gr.Button('Previous');nxt=gr.Button('Next',variant='primary')
 status=gr.Markdown()
 outputs=[state,question,*imgs,*radios,prev,nxt,status]
 demo.load(init,outputs=outputs,api_name=False)
 prev.click(previous,inputs=[state,*radios],outputs=outputs,api_name=False)
 nxt.click(next_question,inputs=[state,*radios],outputs=outputs,api_name=False)
 gr.Markdown('Responses are visible only to the study owner. Please finish in one session; refreshing the page starts a new session.')
if __name__=='__main__':demo.queue(default_concurrency_limit=8).launch()