#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
甲骨文图片服务
端口: 5174

用法:
    python3 serve_image.py

请求示例:
    GET  http://localhost:5174/?name=㐁
    POST http://localhost:5174/  (body: name=㐁 或 {"name": "㐁"})

返回: 对应类别的图片（固定128x128，原图等比缩放居中补白，右下角标注汉字，JPEG格式）
"""

import io
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from PIL import Image, ImageDraw, ImageFont

# ==================== 配置 ====================
# 数据集根目录（map.yaml 所在目录）
DATASET_ROOT = Path('/Users/HonestQiao/Downloads/HUST-OBS/data/HUST-OBC-Selected')
MAP_FILE = DATASET_ROOT / 'map.yaml'
PORT = 5174
OUTPUT_SIZE = 128  # 输出图片固定尺寸 (正方形)
# =============================================


def parse_map_yaml(filepath):
    """手动解析简单的 YAML map 文件"""
    name_to_path = {}
    with open(filepath, 'r', encoding='utf-8') as f:
        in_names = False
        for line in f:
            stripped = line.rstrip('\n\r')
            # 跳过空行和注释
            if not stripped.strip() or stripped.strip().startswith('#'):
                continue
            # 检测 names: 开头
            if stripped.strip() == 'names:':
                in_names = True
                continue
            # 解析键值对（支持 Unicode 汉字作为 key）
            if in_names and ':' in stripped:
                # 找到第一个冒号作为分隔符
                key, value = stripped.split(':', 1)
                key = key.strip()
                value = value.strip()
                if key and value:
                    name_to_path[key] = value
    return name_to_path


def load_map():
    """加载 map.yaml"""
    if not MAP_FILE.exists():
        print(f"Error: map.yaml not found: {MAP_FILE}")
        sys.exit(1)
    return parse_map_yaml(MAP_FILE)


BG_COLOR = (255, 255, 255)  # 补白背景色
FONT_PATH = '/Users/HonestQiao/Downloads/AlibabaPuHuiTi-3/AlibabaPuHuiTi-3-55-Regular/AlibabaPuHuiTi-3-55-Regular.ttf'
FONT_SIZE = 24  # 右下角汉字字体大小
TEXT_COLOR = (255, 0, 0)  # 红色文字，醒目
MARGIN = 8  # 文字距离右下角的边距


def resize_image(img_path, name=None):
    """打开图片，等比缩放后居中补白，右下角标注汉字，返回固定尺寸 JPEG 字节流"""
    with Image.open(img_path) as img:
        # 确保 RGB 模式
        if img.mode != 'RGB':
            img = img.convert('RGB')

        # 等比缩放，最长边不超过 OUTPUT_SIZE
        img.thumbnail((OUTPUT_SIZE, OUTPUT_SIZE), Image.LANCZOS)

        # 创建白色画布
        canvas = Image.new('RGB', (OUTPUT_SIZE, OUTPUT_SIZE), BG_COLOR)

        # 居中粘贴
        paste_x = (OUTPUT_SIZE - img.width) // 2
        paste_y = (OUTPUT_SIZE - img.height) // 2
        canvas.paste(img, (paste_x, paste_y))

        # 右下角绘制汉字
        if name:
            try:
                font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
            except Exception:
                font = ImageFont.load_default()

            draw = ImageDraw.Draw(canvas)
            # 获取文字尺寸（Pillow >= 8.0 使用 getbbox）
            try:
                bbox = draw.textbbox((0, 0), name, font=font)
                text_w = bbox[2] - bbox[0]
                text_h = bbox[3] - bbox[1]
            except Exception:
                text_w, text_h = draw.textsize(name, font=font)

            text_x = OUTPUT_SIZE - text_w - MARGIN
            text_y = OUTPUT_SIZE - text_h - MARGIN
            draw.text((text_x, text_y), name, fill=TEXT_COLOR, font=font)

        # 输出到内存
        buf = io.BytesIO()
        canvas.save(buf, format='JPEG', quality=85)
        return buf.getvalue()


class ImageHandler(BaseHTTPRequestHandler):
    name_map = None

    def log_message(self, format, *args):
        """自定义日志输出"""
        print(f"[{self.address_string()}] {format % args}")

    def _get_name_from_query(self):
        """从 URL query 参数获取 name"""
        parsed = urlparse(self.path)
        query = parse_qs(parsed.query)
        return query.get('name', [None])[0]

    def _get_name_from_post(self):
        """从 POST body 获取 name（支持 form 和 json）"""
        content_length = int(self.headers.get('Content-Length', 0))
        if content_length == 0:
            return None

        body = self.rfile.read(content_length).decode('utf-8')
        content_type = self.headers.get('Content-Type', '').lower()

        # application/json
        if 'application/json' in content_type:
            try:
                data = json.loads(body)
                return data.get('name')
            except json.JSONDecodeError:
                return None

        # application/x-www-form-urlencoded
        if 'application/x-www-form-urlencoded' in content_type or '&' in body or '=' in body:
            params = parse_qs(body)
            return params.get('name', [None])[0]

        # 默认尝试 form 解析
        params = parse_qs(body)
        return params.get('name', [None])[0]

    def _handle_request(self, name):
        """处理请求的核心逻辑，返回 (status, img_bytes)"""
        if not name:
            return 400, b"Missing parameter: name"

        # 查找对应的图片路径
        if name not in self.name_map:
            return 404, b"Category not found"

        rel_path = self.name_map[name]
        img_path = DATASET_ROOT / rel_path

        if not img_path.exists():
            return 404, b"Image file not found"

        try:
            img_data = resize_image(img_path, name=name)
            return 200, img_data
        except Exception as e:
            return 500, f"Image processing failed: {e}".encode('utf-8')

    def _send_response(self, status, data):
        """发送响应"""
        if status != 200:
            self.send_error(status, data.decode('utf-8'))
            return

        self.send_response(200)
        self.send_header('Content-Type', 'image/jpeg')
        self.send_header('Content-Length', str(len(data)))
        self.send_header('Access-Control-Allow-Origin', '*')
        self.end_headers()
        self.wfile.write(data)

    def do_OPTIONS(self):
        """处理 OPTIONS 预检请求（CORS）"""
        self.send_response(204)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type')
        self.end_headers()

    def do_HEAD(self):
        """处理 HEAD 请求"""
        name = self._get_name_from_query()
        status, data = self._handle_request(name)
        self.send_response(status)
        if status == 200:
            self.send_header('Content-Type', 'image/jpeg')
            self.send_header('Content-Length', str(len(data)))
            self.send_header('Access-Control-Allow-Origin', '*')
        self.end_headers()

    def do_GET(self):
        """处理 GET 请求"""
        name = self._get_name_from_query()
        ts = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{ts}] GET  name={name if name else '(missing)'}")
        status, data = self._handle_request(name)
        self._send_response(status, data)

    def do_POST(self):
        """处理 POST 请求"""
        name = self._get_name_from_post()
        ts = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{ts}] POST name={name if name else '(missing)'}")
        status, data = self._handle_request(name)
        self._send_response(status, data)


def run_server():
    name_map = load_map()
    print(f"已加载 {len(name_map)} 个类别映射")
    for name, path in name_map.items():
        print(f"  {name} -> {path}")

    ImageHandler.name_map = name_map

    server = HTTPServer(('0.0.0.0', PORT), ImageHandler)
    print(f"\n服务已启动: http://localhost:{PORT}/")
    print(f'GET 示例: curl "http://localhost:{PORT}/?name=㐁" -o output.jpg')
    print(f'POST示例: curl -X POST -d "name=㐁" "http://localhost:{PORT}/" -o output.jpg')
    print(f"输出图片: 固定 {OUTPUT_SIZE}x{OUTPUT_SIZE}，原图等比缩放居中，右下角标注红色汉字，JPEG 格式")
    print("按 Ctrl+C 停止\n")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n服务已停止")
        server.shutdown()


if __name__ == '__main__':
    run_server()
