#!/usr/bin/env python3
"""Python 3.10+, standard library. AS Gemini tool roundtrip or native MP4 understanding."""
import argparse
import base64
import json
import os
import re
import urllib.request
from pathlib import Path


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl): return None


def post(key, path, body):
    req = urllib.request.Request('https://openpix.cc' + path,
        data=json.dumps(body, ensure_ascii=False).encode('utf8'),
        headers={'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'}, method='POST')
    # Billable POST, deliberately no automatic retry.
    with urllib.request.build_opener(NoRedirect()).open(req, timeout=180) as response:
        value = json.load(response)
    if value.get('error'): raise RuntimeError(json.dumps(value['error'], ensure_ascii=False))
    return value


def run_tools(key, model):
    tools = [{'type': 'function', 'function': {
        'name': 'lookup_stock', 'description': '查询指定商品和仓库的库存',
        'parameters': {'type': 'object', 'properties': {'sku': {'type': 'string'}, 'warehouse': {'type': 'string'}},
                       'required': ['sku', 'warehouse']}}}]
    messages = [{'role': 'user', 'content': '请查询 TEA-01 在上海仓的库存，告诉我能否发货。'}]
    first = post(key, '/v1/chat/completions', {'model': model, 'messages': messages, 'tools': tools, 'tool_choice': 'auto', 'stream': False})
    if not first.get('choices'): raise RuntimeError('No assistant choice')
    assistant = first['choices'][0]['message']
    calls = assistant.get('tool_calls') or []
    if not calls: raise RuntimeError('No tool call; inspect answer instead of assuming tool success')
    messages.append(assistant)  # Keep COMPLETE assistant message and tool metadata.
    for call in calls:
        function = call['function']
        if function['name'] != 'lookup_stock' or not call.get('id'): raise RuntimeError('Unexpected tool')
        args = json.loads(function['arguments'])
        if not isinstance(args, dict) or not all(isinstance(args.get(k), str) for k in ('sku', 'warehouse')):
            raise RuntimeError('Invalid tool arguments')
        # Fixture only: replace this explicitly allowed function with your application lookup.
        result = {'sku': args['sku'], 'warehouse': args['warehouse'],
                  'stock': 17 if args['sku'] == 'TEA-01' and args['warehouse'] in ('上海', '上海仓') else 0,
                  'source': 'demo_fixture'}
        messages.append({'role': 'tool', 'tool_call_id': call['id'], 'content': json.dumps(result, ensure_ascii=False)})
    second = post(key, '/v1/chat/completions', {'model': model, 'messages': messages, 'tools': tools, 'tool_choice': 'auto', 'stream': False})
    print(json.dumps({'first_usage': first.get('usage'), 'second_response': second}, ensure_ascii=False, indent=2))


def video_body(file, prompt):
    path = Path(file)
    # Conservative sample limit, not the service's maximum.
    if path.suffix.lower() != '.mp4' or not 0 < path.stat().st_size <= 16 * 1024 * 1024:
        raise ValueError('Example accepts MP4 up to 16 MiB; start with a short clip')
    raw = path.read_bytes()
    if len(raw) < 12 or raw[4:8] != b'ftyp': raise ValueError('Expected MP4 ftyp signature')
    return {'contents': [{'role': 'user', 'parts': [
        {'inlineData': {'mimeType': 'video/mp4', 'data': base64.b64encode(raw).decode('ascii')}}, {'text': prompt}]}]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('mode', choices=['tools', 'video'])
    parser.add_argument('--model', default='gemini-3.8-flash')
    parser.add_argument('--file')
    parser.add_argument('--prompt', default='按时间顺序描述视频中的主体、动作和文字，不猜测不可见内容。')
    args = parser.parse_args();key = os.environ.get('MXB_API_KEY')
    if not key: parser.error('Set MXB_API_KEY')
    if not re.fullmatch(r'[A-Za-z0-9._-]+', args.model): parser.error('Invalid model ID')
    if args.mode == 'tools': run_tools(key, args.model)
    else:
        if not args.file: parser.error('--file is required for video')
        response = post(key, '/v1beta/models/' + args.model + ':generateContent', video_body(args.file, args.prompt))
        texts = [p['text'] for c in response.get('candidates', []) for p in c.get('content', {}).get('parts', []) if isinstance(p.get('text'), str)]
        if not texts: raise RuntimeError('No text; inspect finishReason/promptFeedback before another POST')
        print('\n'.join(texts));print(json.dumps({'usageMetadata': response.get('usageMetadata')}, ensure_ascii=False, indent=2))


if __name__ == '__main__': main()