unified_server.py

"""
统一服务器 - 合并H5页面 + OCR + 导出服务于一个端口
手机端通过内网穿透访问
"""
import json
import io
import os
import sys
import re
import base64
import tempfile
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse

# 确保能导入同目录模块
WORK_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, WORK_DIR)

# 导入导出服务核心函数
from export_server import generate_docx

# 导入OCR核心函数
import pytesseract
import platform
import shutil
from PIL import Image

# 跨平台 Tesseract 路径配置
if platform.system() == 'Windows':
    # Windows 本地开发环境
    pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
    os.environ['TESSDATA_PREFIX'] = os.path.expandvars(r'%USERPROFILE%\.tesseract\tessdata')
else:
    # Linux/Docker 环境：tesseract 已通过 apt 安装到 PATH
    tesseract_bin = shutil.which('tesseract')
    if tesseract_bin:
        pytesseract.pytesseract.tesseract_cmd = tesseract_bin
    # 自动检测 tessdata 路径（兼容 Tesseract 4.x/5.x）
    for _tessdata in ['/usr/share/tesseract-ocr/5/tessdata',
                      '/usr/share/tesseract-ocr/4.00/tessdata',
                      '/usr/share/tessdata']:
        if os.path.isdir(_tessdata):
            os.environ.setdefault('TESSDATA_PREFIX', _tessdata)
            break
    print(f'[Tesseract] cmd={pytesseract.pytesseract.tesseract_cmd}, '
          f'TESSDATA_PREFIX={os.environ.get("TESSDATA_PREFIX", "(default)")}')

# H5页面路径
H5_FILE = os.path.join(WORK_DIR, 'test-h5.html')

PORT = int(os.environ.get('PORT', 7860))


def ocr_image(image_data):
    """对图片进行OCR识别"""
    img = Image.open(io.BytesIO(image_data))
    w, h = img.size
    max_dim = 2000
    if w > max_dim or h > max_dim:
        ratio = min(max_dim / w, max_dim / h)
        img = img.resize((int(w * ratio), int(h * ratio)), Image.LANCZOS)
    text = pytesseract.image_to_string(img, lang='chi_sim', config='--psm 6')
    return text.strip()


class UnifiedHandler(BaseHTTPRequestHandler):
    
    def _add_cors_headers(self):
        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')

    def do_OPTIONS(self):
        self.send_response(200)
        self._add_cors_headers()
        self.end_headers()

    def do_GET(self):
        """处理GET请求 - 返回H5页面"""
        path = urlparse(self.path).path
        
        if path == '/' or path == '/test-h5.html' or path == '/index.html':
            try:
                with open(H5_FILE, 'r', encoding='utf-8') as f:
                    content = f.read()
                self.send_response(200)
                self.send_header('Content-Type', 'text/html; charset=utf-8')
                self._add_cors_headers()
                self.end_headers()
                self.wfile.write(content.encode('utf-8'))
            except Exception as e:
                self.send_error(500, f'读取页面失败: {e}')
        elif path == '/health':
            # 健康检查
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self._add_cors_headers()
            self.end_headers()
            self.wfile.write(json.dumps({'status': 'ok'}).encode('utf-8'))
        else:
            self.send_error(404, '页面不存在')

    def do_POST(self):
        """处理POST请求 - OCR和导出"""
        path = urlparse(self.path).path
        
        if path == '/ocr':
            self._handle_ocr()
        elif path == '/export':
            self._handle_export()
        else:
            self.send_error(404, '接口不存在')

    def _handle_ocr(self):
        """处理OCR识别请求"""
        content_length = int(self.headers.get('Content-Length', 0))
        body = self.rfile.read(content_length)
        
        content_type = self.headers.get('Content-Type', '')
        if 'multipart/form-data' not in content_type:
            self._send_json({'error': '需要 multipart/form-data'}, 400)
            return
        
        # 提取 boundary
        boundary = content_type.split('boundary=')[1].strip()
        if boundary.startswith('"') and boundary.endswith('"'):
            boundary = boundary[1:-1]
        
        # 手动解析 multipart
        parts = body.split(('--' + boundary).encode())
        image_data = None
        for part in parts:
            if b'Content-Disposition' in part and b'filename=' in part:
                header_end = part.find(b'\r\n\r\n')
                if header_end == -1:
                    continue
                file_data = part[header_end + 4:]
                if file_data.endswith(b'\r\n'):
                    file_data = file_data[:-2]
                image_data = file_data
                break
        
        if not image_data:
            self._send_json({'error': '未找到图片数据'}, 400)
            return
        
        try:
            text = ocr_image(image_data)
            self._send_json({'text': text, 'length': len(text)})
            print(f'[OCR] 识别成功，提取 {len(text)} 个字符')
        except Exception as e:
            self._send_json({'error': f'OCR识别失败: {str(e)}'}, 500)
            print(f'[OCR] 识别失败: {e}')

    def _handle_export(self):
        """处理Word文档导出请求"""
        content_length = int(self.headers.get('Content-Length', 0))
        body = self.rfile.read(content_length)
        
        try:
            visit_data = json.loads(body)
            buffer = generate_docx(visit_data)
            
            self.send_response(200)
            self.send_header('Content-Type',
                             'application/vnd.openxmlformats-officedocument.wordprocessingml.document')
            self._add_cors_headers()
            filename = 'visit_%s_%sstores.docx' % (
                visit_data.get('visitDateDisplay', '').replace('.', '-'),
                visit_data.get('totalCount', 0)
            )
            self.send_header('Content-Disposition', f'attachment; filename="{filename}"')
            self.end_headers()
            self.wfile.write(buffer.getvalue())
            print(f'[Export] 导出成功: {filename}')
        except Exception as e:
            self.send_response(500)
            self.send_header('Content-Type', 'application/json')
            self._add_cors_headers()
            self.end_headers()
            self.wfile.write(json.dumps({'error': str(e)}).encode('utf-8'))
            print(f'[Export] 导出失败: {e}')

    def _send_json(self, data, status=200):
        self.send_response(status)
        self.send_header('Content-Type', 'application/json; charset=utf-8')
        self._add_cors_headers()
        self.end_headers()
        self.wfile.write(json.dumps(data, ensure_ascii=False).encode('utf-8'))

    def log_message(self, format, *args):
        print(f'[{self.log_date_time_string()}] {args[0]}')


if __name__ == '__main__':
    print('=' * 50)
    print('  药店拜访助手 - 统一服务器')
    print(f'  本机访问: http://127.0.0.1:{PORT}')
    print(f'  局域网访问: http://<本机IP>:{PORT}')
    print('=' * 50)
    
    # 绑定 0.0.0.0 让局域网设备也能访问
    server = HTTPServer(('0.0.0.0', PORT), UnifiedHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print('\n服务器已停止')
        server.shutdown()