#!/usr/bin/env python3
"""Wiki Chat Proxy - Proxies chat requests to z.ai GLM-4-Plus"""

import json
import os
import urllib.request

ZAI_TOKEN = os.environ.get('ZAI_API_TOKEN', '')

def handler(event):
    try:
        if event.get('method') != 'POST':
            return {'status': 405, 'body': 'Method not allowed'}
        
        body = json.loads(event.get('body', '{}'))
        messages = body.get('messages', [])
        
        if not messages:
            return {'status': 400, 'body': json.dumps({'error': 'messages required'})}
        
        api_url = 'https://api.z.ai/api/coding/paas/v4/chat/completions'
        
        payload = json.dumps({
            'model': 'glm-4-plus',
            'messages': messages,
            'temperature': 0.7,
            'max_tokens': 2000,
        }).encode()
        
        headers = {'Content-Type': 'application/json'}
        if ZAI_TOKEN:
            headers['Authorization'] = 'Bearer ' + ZAI_TOKEN
        
        req = urllib.request.Request(api_url, data=payload, headers=headers, method='POST')
        
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                data = json.loads(resp.read().decode())
                content = data.get('choices', [{}])[0].get('message', {}).get('content', '')
                return {
                    'status': 200,
                    'headers': {'Content-Type': 'application/json'},
                    'body': json.dumps({'response': content})
                }
        except urllib.error.HTTPError as e:
            err_body = e.read().decode() if e.fp else ''
            try:
                err_json = json.loads(err_body)
                err_msg = err_json.get('error', {}).get('message', err_json.get('message', str(e)))
            except:
                err_msg = str(e)
            return {'status': e.code, 'body': json.dumps({'error': err_msg})}
    except Exception as e:
        return {'status': 500, 'body': json.dumps({'error': str(e)})}