import base64,io,os,re,tempfile,time
import gradio as gr
import requests
from PIL import Image,ImageDraw
CHAT='openai/gpt-4.1'; PLAN='gemini-3.1-flash-lite'; GEMINI_IMAGE='gemini-2.5-flash-image'; ATTEMPTS=5
CHAT_SYS='''You are Pixelship, a concise creative chatbot. Chat normally. Only when asked to create an image, end with (TOOL:TXTTOIMGOUT, prompt:DETAILED PROMPT). If an image is attached or an edit is requested, use IMGTOIMGOUT. Always specify a complete background. Never invent brands; user-mentioned brand/model names must be inside double quotes.'''
BLUE_SYS='''Make a strict 16x16 blueprint with a full background, but make the composition creative: angled subjects, diagonals, asymmetry, foreground/background layers, perspective cues, and expressive staging are allowed. If text is needed, use TEXT() commands as semantic exact words, not as a requirement for pixel-art font. For posters/flyers/ads/thumbnails/covers, include every requested text line and make the layout richly designed with tasteful clutter: badges, panels, icons, arrows, callouts, feature boxes, stickers, borders, and layered elements. Output ROWS:16, PIXELS_PER_ROW:16, PALETTE entries 1-16 as six-digit hex, CREATIVITY, and ROW1-ROW16 with 16 comma-separated palette numbers or TEXT() values. Output only the blueprint.'''
QUALITY='''IMAGE 1 is the candidate; later images are references. Score 1-10. PASS requires 7+, correct subject/reference identity, composition, background, requested text, and no gibberish. If uploaded/user references and web references are both supplied, uploaded/user references are primary and web references are supplemental only. The candidate must preserve uploaded/user reference identity first; web refs may add current/factual details but must not override the uploaded refs. Ignore microscopic defects. Reply exactly (img:PASS score:N correction:) or (img:FALSE score:N correction:reason).'''
MERCH=re.compile(r'\b(merch|merchandise|t-?shirt|tee|hoodie|sweatshirt|poster|art print|sticker|decal|mug|phone case|case cover|skin|costume|plush|plushie|figurine|keychain|lanyard|wall art)\b',re.I)
def run(token,model,data):
 h={'Authorization':'Bearer '+token.strip(),'Content-Type':'application/json'}
 if not token.strip(): raise gr.Error('Enter your Replicate API token.')
 for wait in (0,1,2,4,8):
  if wait: time.sleep(wait)
  try:
   r=requests.post('https://api.replicate.com/v1/models/'+model+'/predictions',headers=h,json={'input':data},timeout=60)
   if r.status_code in (408,429,500,502,503,504): continue
   r.raise_for_status(); p=r.json(); break
  except requests.RequestException:
   if wait==8: raise
 while p['status'] not in ('succeeded','failed','canceled','aborted'):
  time.sleep(1); p=requests.get(p['urls']['get'],headers=h,timeout=60).json()
 if p['status']!='succeeded': raise gr.Error(str(p.get('error') or 'Replicate failed'))
 o=p.get('output'); return (o[0] if len(o)==1 else ''.join(map(str,o))) if isinstance(o,list) else (o.get('url') if isinstance(o,dict) else str(o))
def gemini_part(source):
 if str(source).startswith('data:'):
  m=re.match(r'^data:([^;]+);base64,(.+)$',source)
  if not m: raise gr.Error('Gemini image input must be a data URL.')
  return {'inline_data':{'mime_type':m.group(1),'data':m.group(2)}}
 r=requests.get(source,headers={'User-Agent':'Mozilla/5.0'},timeout=60)
 r.raise_for_status()
 return {'inline_data':{'mime_type':r.headers.get('content-type') or 'image/jpeg','data':base64.b64encode(r.content).decode()}}
def gemini_text(key,model,data):
 if not key.strip(): raise gr.Error('Enter your Gemini API key.')
 parts=[{'text':str(data.get('prompt',''))}]
 for image in (data.get('image_input') or [])[:16]: parts.append(gemini_part(image))
 body={'contents':[{'role':'user','parts':parts}]}
 if data.get('system_prompt'): body['systemInstruction']={'parts':[{'text':str(data.get('system_prompt'))}]}
 url='https://generativelanguage.googleapis.com/v1beta/models/'+model+':generateContent?key='+key.strip()
 last=None
 for wait in (0,1,2,4,8,16):
  if wait: time.sleep(wait)
  r=requests.post(url,json=body,timeout=120)
  if r.status_code in (408,429,500,502,503,504): last=r.text; continue
  try: r.raise_for_status()
  except requests.HTTPError as e: raise gr.Error(r.text[:500]) from e
  j=r.json(); break
 else:
  raise gr.Error(last or 'Gemini text request failed')
 text=''.join(part.get('text','') for part in j.get('candidates',[{}])[0].get('content',{}).get('parts',[])).strip()
 if not text: raise gr.Error('Gemini text response was empty.')
 return text
def gemini_image(key,prompt,image_data_url,aspect='match_input_image'):
 if not key.strip(): raise gr.Error('Enter your Gemini API key.')
 parts=[{'text':prompt+'\nRender one polished image. Preferred aspect ratio: '+aspect+'.'}]
 if image_data_url:
  m=re.match(r'^data:([^;]+);base64,(.+)$',image_data_url)
  if not m: raise gr.Error('Gemini image input must be a data URL.')
  parts.append({'inline_data':{'mime_type':m.group(1),'data':m.group(2)}})
 body={'contents':[{'role':'user','parts':parts}],'generationConfig':{'responseModalities':['IMAGE']}}
 url='https://generativelanguage.googleapis.com/v1beta/models/'+GEMINI_IMAGE+':generateContent?key='+key.strip()
 last=None
 for wait in (0,1,2,4,8,16):
  if wait: time.sleep(wait)
  r=requests.post(url,json=body,timeout=120)
  if r.status_code in (408,429,500,502,503,504): last=r.text; continue
  try: r.raise_for_status()
  except requests.HTTPError as e: raise gr.Error(r.text[:500]) from e
  j=r.json(); break
 else:
  raise gr.Error(last or 'Gemini image request failed')
 for part in j.get('candidates',[{}])[0].get('content',{}).get('parts',[]):
  blob=part.get('inlineData') or part.get('inline_data') or {}
  data=blob.get('data'); mime=blob.get('mimeType') or blob.get('mime_type') or 'image/jpeg'
  if data: return 'data:'+mime+';base64,'+data
 raise gr.Error('Gemini did not return an image.')
def data_url(im):
 im=im.convert('RGB'); im.thumbnail((768,768)); b=io.BytesIO(); im.save(b,'JPEG',quality=80,optimize=True); return 'data:image/jpeg;base64,'+base64.b64encode(b.getvalue()).decode()
def splitrow(s):
 out=[]; cur=''; depth=0
 for c in s:
  depth+=c=='('; depth-=c==')'
  if c==',' and depth==0: out.append(cur.strip()); cur=''
  else: cur+=c
 if cur.strip(): out.append(cur.strip())
 return out
def blueprint(text):
 pal={int(n):c for n,c in re.findall(r'^\s*(1[0-6]|[1-9])\s*:\s*#?([0-9a-fA-F]{6})\s*$',text,re.M)}
 if len(pal)!=16: raise gr.Error('Blueprint palette was incomplete.')
 rows=[]
 for i in range(1,17):
  m=re.search(r'^\s*ROW'+str(i)+r'\s*:\s*(?:\((.*?)\)|(.*?))\s*$',text,re.M); cells=splitrow((m.group(1) or m.group(2)).strip()) if m else []
  vals=[int(x) if x.isdigit() and int(x) in pal else 1 for x in cells[:16]]; vals += [vals[-1] if vals else 1]*(16-len(vals)); rows.append(vals)
 im=Image.new('RGB',(256,256)); d=ImageDraw.Draw(im)
 for y,row in enumerate(rows):
  for x,v in enumerate(row): d.rectangle((x*16,y*16,x*16+15,y*16+15),fill='#'+pal[v])
 return im
def collage(ref,bp):
 largest=max(ref.width,ref.height); scale=1024/largest if largest>1024 else (512/largest if largest<512 else 1.0); w=max(1,round(ref.width*scale)); h=max(1,round(ref.height*scale))
 c=Image.new('RGB',(w,h),'#0f1116'); main=ref.copy(); main.thumbnail((w,h)); c.paste(main,((w-main.width)//2,(h-main.height)//2))
 inset=max(24,min(round(w*.32),round(h*.42),320)); panel=bp.copy().resize((inset,inset),Image.Resampling.NEAREST)
 border=max(4,round(inset*.025)); framed=Image.new('RGB',(inset+border*2,inset+border*2),'#b8ff4d'); framed.paste(panel,(border,border))
 c.paste(framed,(w-framed.width-border,h-framed.height-border))
 return c
def search_queries(token,prompt):
 raw=run(token,CHAT,{'prompt':'Extract 1 to 3 web search queries for this image request. Search the subject/design, not the requested output format. Example: ios 26 poster with all new features -> ios 26, ios 26 features. Remove poster/wallpaper/banner/render/style words unless they are the actual subject. Return comma-separated queries only. Request: '+prompt,'system_prompt':'You extract concise search queries.'})
 return [q.strip() for q in re.split(r'[,\n]+',raw) if q.strip()][:3] or [prompt]
def serper_web(key,query):
 r=requests.post('https://google.serper.dev/search',headers={'X-API-KEY':key.strip(),'Content-Type':'application/json'},json={'q':query,'num':5},timeout=20)
 r.raise_for_status(); return '\n'.join((i.get('title','')+': '+i.get('snippet','')).strip() for i in (r.json().get('organic') or [])[:5])
def collage_many(refs,bp):
 if not refs: return bp
 thumbs=[]
 for ref in refs[:6]:
  im=ref.copy(); im.thumbnail((512,512)); thumbs.append(im)
 cols=2 if len(thumbs)>1 else 1; rows=(len(thumbs)+cols-1)//cols
 cell=512; c=Image.new('RGB',(cols*cell,rows*cell),'#0f1116')
 for i,im in enumerate(thumbs):
  x=(i%cols)*cell+(cell-im.width)//2; y=(i//cols)*cell+(cell-im.height)//2; c.paste(im,(x,y))
 inset=220; panel=bp.copy().resize((inset,inset),Image.Resampling.NEAREST); framed=Image.new('RGB',(inset+12,inset+12),'#b8ff4d'); framed.paste(panel,(6,6)); c.paste(framed,(c.width-framed.width-12,c.height-framed.height-12))
 return c
def serper_reference(key,query):
 if not key.strip(): raise gr.Error('Enter your Serper API key or turn web search off.')
 r=requests.post('https://google.serper.dev/images',headers={'X-API-KEY':key.strip(),'Content-Type':'application/json'},json={'q':query,'num':10},timeout=20)
 r.raise_for_status(); results=r.json().get('images') or []
 for item in results:
  title=(item.get('title') or '')+' '+(item.get('source') or '')+' '+(item.get('link') or '')
  src=item.get('imageUrl') or item.get('thumbnailUrl')
  if not src or MERCH.search(title): continue
  try:
   img=requests.get(src,headers={'User-Agent':'Mozilla/5.0'},timeout=20).content
   im=Image.open(io.BytesIO(img)).convert('RGB')
   if im.width<80 or im.height<80: continue
   return im,src
  except Exception:
   continue
 raise gr.Error('Web search found no usable non-merch images.')
def judge(gemini_key,prompt,result,refs):
 v=gemini_text(gemini_key,PLAN,{'prompt':'Request: '+prompt+'\nJudge IMAGE 1 strictly.','system_prompt':QUALITY,'image_input':[result]+refs}); m=re.search(r'score:\s*(10|[1-9])',v,re.I); score=int(m.group(1)) if m else 0; corr=(re.search(r'correction:([^)]*)',v,re.I) or [0,'Improve accuracy'])[1]; return bool(re.search(r'img:PASS',v,re.I)) and score>=7,score,corr
def generate(token,gemini_key,prompt,refpath=None,web_key='',web_search=False):
 ref=None; ref_imgs=[]; user_refs=[]; web_refs=[]; refs=[]; note=''; research=''
 if refpath:
  ref=Image.open(refpath).convert('RGB'); user_refs=[ref]; ref_imgs=[ref]; note='uploaded reference'
 if web_search:
  queries=search_queries(token,prompt); research='\n'.join(serper_web(web_key,q) for q in queries[:3])
  sources=[]
  for q in queries[:3]:
   try:
    im,src=serper_reference(web_key,q); web_refs.append(im); sources.append(src)
   except Exception: pass
  if web_refs and not ref: ref=web_refs[0]
  ref_imgs=user_refs+web_refs
  note=(note+' + supplemental web references: ' if note else 'web references: ')+', '.join(sources[:3])
 pi={'prompt':prompt,'system_prompt':BLUE_SYS}
 if user_refs and web_refs:
  pi['prompt']+='''\nUploaded/user reference images are PRIMARY. Web images are supplemental only for missing/current details. Do not let web images override, replace, dominate, or distract from the uploaded reference.'''
 if ref_imgs: pi['image_input']=[data_url(x) for x in ref_imgs[:6]]
 bp=blueprint(gemini_text(gemini_key,PLAN,pi)); inp=collage_many(ref_imgs,bp) if ref_imgs else bp; refs=[data_url(x) for x in ref_imgs[:6]]
 base='Create a polished creative image from the supplied blueprint/reference. The blueprint is a composition guide, not a pixel-art font/style requirement. Allow dynamic angles, depth, layered foreground/background, expressive staging, and tasteful visual density. For posters/flyers/ads/thumbnails/covers, make the design rich and information-dense with readable panels, badges, icons, callouts, and decorative clutter. '+prompt+'. Preserve reference identity and physical design. Use a complete background and no gibberish. Render any requested text exactly with an appropriate font, not automatically a pixel font.'
 if note: base+=' The uploaded/user reference is the primary visual source. If web references are present, use them only as supplemental context for missing/current design details; never let web results override, replace, dominate, or distract from uploaded/user references. Do not copy a poster/layout from search results.'
 if research: base+=' Web research notes for factual feature/design details, not layout copying:\n'+research
 current=base; candidates=[]
 for _ in range(ATTEMPTS):
  image_data=data_url(inp)
  result=gemini_image(gemini_key,current,image_data,'match_input_image' if ref else '1:1'); passed,score,corr=judge(gemini_key,prompt,result,refs); candidates.append((score,result))
  if passed: return result
  current=base+' Improve the result while keeping the same reference and blueprint. Fix: '+corr
 return max(candidates)[1]
def download(u):
 fd,p=tempfile.mkstemp(suffix='.png'); os.close(fd)
 if str(u).startswith('data:image/'):
  open(p,'wb').write(base64.b64decode(str(u).split(',',1)[1])); return p
 r=requests.get(u,timeout=90); r.raise_for_status(); open(p,'wb').write(r.content); return p
def send(msg,history,token,gemini_key,attachment,web_key,web_search):
 if not msg.strip(): return history,'',attachment
 transcript='\n'.join(x['role']+': '+x['content'] for x in history if isinstance(x['content'],str)); reply=run(token,CHAT,{'prompt':transcript+'\nuser: '+msg+'\nassistant:','system_prompt':CHAT_SYS+('\nAn image is attached.' if attachment else '')+('\nThe user enabled web search. If creating an image, use the searched image as a visual reference.' if web_search else '')}); m=re.search(r'\(TOOL:(TXTTOIMGOUT|IMGTOIMGOUT)\s*,\s*prompt:(.*)\)\s*$',reply,re.I|re.S); history=history+[{'role':'user','content':msg}]
 if not m: return history+[{'role':'assistant','content':reply}],'',attachment
 visible=reply[:m.start()].strip() or 'Here is what I made.'; result=generate(token,gemini_key,m.group(2).strip(),attachment if m.group(1).upper()=='IMGTOIMGOUT' else None,web_key,web_search); history += [{'role':'assistant','content':visible},{'role':'assistant','content':gr.Image(value=download(result))}]; return history,'',None
with gr.Blocks(title='Pixelship Gradio') as demo:
 gr.Markdown('# Pixelship\nStandalone chatbot image pipeline')
 token=gr.Textbox(label='Replicate API token (chat only)',type='password',placeholder='r8_...')
 gemini_key=gr.Textbox(label='Gemini API key (blueprint, judging, image rendering)',type='password',placeholder='AIza...')
 web_key=gr.Textbox(label='Serper API key (optional, for web search references)',type='password',placeholder='serper.dev API key')
 web_search=gr.Checkbox(label='Use web search reference',value=False)
 chat=gr.Chatbot(height=600)
 attach=gr.Image(label='Optional edit reference',type='filepath')
 msg=gr.Textbox(label='Message',lines=2)
 button=gr.Button('Send',variant='primary')
 button.click(send,[msg,chat,token,gemini_key,attach,web_key,web_search],[chat,msg,attach])
 msg.submit(send,[msg,chat,token,gemini_key,attach,web_key,web_search],[chat,msg,attach])
if __name__=='__main__': demo.queue().launch()