"""ROCm/CUDA local WebUI for zero-shot vision transfer. No remote requests."""
import argparse,base64,io,json,sys,threading,traceback
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
from PIL import Image
from mugi_decision.vision import VisionDecisionEngine

ROOT=Path(__file__).resolve().parents[1]

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--model',default=str(ROOT))
    ap.add_argument('--adapter')
    ap.add_argument('--port',type=int,default=7861)
    ap.add_argument('--optimization',choices=['auto','none','hip','fast'],default='auto')
    args=ap.parse_args()
    engine=VisionDecisionEngine(model_path=args.model,adapter=args.adapter,optimization=args.optimization)
    lock=threading.Lock()
    class Handler(BaseHTTPRequestHandler):
        def reply(self,status,data,kind='application/json; charset=utf-8'):
            raw=json.dumps(data,ensure_ascii=False).encode('utf-8') if kind.startswith('application/json') else data
            self.send_response(status)
            self.send_header('Content-Type',kind)
            self.send_header('Content-Length',str(len(raw)))
            self.send_header('Cache-Control','no-store')
            self.send_header('X-Content-Type-Options','nosniff')
            self.end_headers();self.wfile.write(raw)
        def do_GET(self):
            if self.path=='/':self.reply(200,(ROOT/'webui/vision.html').read_bytes(),'text/html; charset=utf-8')
            elif self.path=='/api/status':self.reply(200,{'ready':True,**engine.metadata})
            else:self.reply(404,{'error':'Not found'})
        def do_POST(self):
            if self.path!='/api/predict':return self.reply(404,{'error':'Not found'})
            origin=self.headers.get('Origin')
            if origin and origin not in (f'http://127.0.0.1:{args.port}',f'http://localhost:{args.port}'):
                return self.reply(403,{'error':'Local origin required'})
            try:
                size=int(self.headers.get('Content-Length','0'))
                if not 0<size<=24*1024*1024:raise ValueError('请求过大（上限24MB）')
                payload=json.loads(self.rfile.read(size))
                picture=None
                if payload.get('image'):
                    encoded=payload['image']
                    if not isinstance(encoded,str) or not encoded.startswith('data:image/'):
                        raise ValueError('需要图片 data URL')
                    picture=Image.open(io.BytesIO(base64.b64decode(encoded.split(',',1)[1],validate=True)))
                    if picture.width*picture.height>20_000_000:raise ValueError('图片不得超过2000万像素')
                    picture.load()
                if not lock.acquire(blocking=False):return self.reply(409,{'error':'模型正在处理另一个请求，请稍后再试'})
                try:
                    result=engine.predict(payload.get('description',''),payload.get('options',[]),picture,
                        int(payload.get('image_size',512)),int(payload.get('max_length',4096)),bool(payload.get('verify_cache',False)))
                finally:lock.release()
                self.reply(200,result)
            except (ValueError,TypeError,KeyError,Image.DecompressionBombError) as exc:
                self.reply(400,{'error':str(exc)})
            except Exception as exc:
                traceback.print_exc()
                self.reply(500,{'error':f'{type(exc).__name__}: {exc}'})
    server=ThreadingHTTPServer(('127.0.0.1',args.port),Handler)
    print(json.dumps(engine.metadata,ensure_ascii=False),flush=True)
    print(f'WebUI ready: http://127.0.0.1:{args.port}',flush=True)
    server.serve_forever()

if __name__=='__main__':main()