#!/usr/bin/env python3
"""Wiki Chat LLM Proxy Server on port 8098 using Flask"""

import json
import os
import urllib.request
from flask import Flask, request, jsonify, make_response

app = Flask(__name__)

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

@app.after_request
def add_cors(response):
    h = 'Access-Control-Allow-Origin'
    response.headers[h] = '*'
    response.headers['Access-Control-Allow-Methods'] = 'POST, OPTIONS'
    response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
    return response

def call_zai(messages, token=None):
    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 token:
        headers['Authorization'] = 'Bearer ' + token
    req = urllib.request.Request(api_url, data=payload, headers=headers, method='POST')
    with urllib.request.urlopen(req, timeout=60) as resp:
        data = json.loads(resp.read().decode())
        return data.get('choices', [{}])[0].get('message', {}).get('content', '')

@app.route('/chat/wiki', methods=['POST', 'OPTIONS'])
def chat_wiki():
    if request.method == 'OPTIONS':
        return make_response('', 200)
    try:
        body = request.get_json(force=True)
        messages = body.get('messages', [])
        if not messages:
            return jsonify({'error': 'messages required'}), 400
        token = body.get('token') or ZAI_TOKEN
        if not token:
            return jsonify({'error': 'No token provided'}), 401
        content = call_zai(messages, token)
        return jsonify({'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 Exception:
            err_msg = str(e)
        return jsonify({'error': err_msg}), e.code
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/chat/wiki-tunnel', methods=['POST', 'OPTIONS'])
def chat_wiki_tunnel():
    if request.method == 'OPTIONS':
        return make_response('', 200)
    try:
        body = request.get_json(force=True)
        messages = body.get('messages', [])
        if not messages:
            return jsonify({'error': 'messages required'}), 400
        token = ZAI_TOKEN
        if not token:
            return jsonify({'error': 'No server token configured. Use Token mode and paste your API key.'}), 503
        content = call_zai(messages, token)
        return jsonify({'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 Exception:
            err_msg = str(e)
        return jsonify({'error': err_msg}), e.code
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    port = int(os.environ.get('WIKI_CHAT_PORT', 8098))
    print(f'[WikiChat] LLM proxy running on port {port}')
    app.run(host='0.0.0.0', port=port)