TorqueWrench/backend/database.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库操作模块
使用SQLite存储工单数据
"""

import sqlite3
import json
from datetime import datetime, timedelta
from typing import Dict, List, Optional


class Database:
    """数据库操作类"""
    
    def __init__(self, db_path: str = "wrench.db"):
        """
        初始化数据库
        :param db_path: 数据库文件路径
        """
        self.db_path = db_path
        self.init_database()
    
    def get_connection(self):
        """获取数据库连接"""
        conn = sqlite3.connect(self.db_path, timeout=30.0, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute('PRAGMA journal_mode=WAL')
        conn.execute('PRAGMA synchronous=NORMAL')
        conn.execute('PRAGMA cache_size=-64000')
        return conn
    
    def init_database(self):
        """初始化数据库表"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # 创建系统配置表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS system_config (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 记录首次启动时间
        cursor.execute("SELECT value FROM system_config WHERE key = 'first_start_time'")
        if not cursor.fetchone():
            cursor.execute("INSERT INTO system_config (key, value) VALUES ('first_start_time', ?)",
                         (datetime.now().strftime("%Y-%m-%d %H:%M:%S"),))
        
        # 创建工单表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS work_orders (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                trace_id TEXT NOT NULL,
                process_id TEXT NOT NULL,
                process_name TEXT,
                product_name TEXT,
                operator TEXT,
                status TEXT DEFAULT 'pending',
                bolts TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(trace_id, process_id)
            )
        ''')
        
        # 创建扳手设备表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS wrench_devices (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                device_name TEXT NOT NULL,
                device_sn TEXT UNIQUE,
                ip_address TEXT NOT NULL,
                port INTEGER DEFAULT 7888,
                address_code INTEGER DEFAULT 1,
                status TEXT DEFAULT 'offline',
                last_heartbeat TIMESTAMP,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建认领表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS claimed_orders (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                trace_id TEXT NOT NULL,
                process_id TEXT NOT NULL,
                process_name TEXT,
                operator TEXT,
                status TEXT DEFAULT 'claimed',
                claimed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                UNIQUE(trace_id, process_id)
            )
        ''')
        
        # 创建结果表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS work_results (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                trace_id TEXT NOT NULL,
                process_id TEXT NOT NULL,
                process_name TEXT,
                bolts TEXT NOT NULL,
                device_sn TEXT,
                device_name TEXT,
                submitted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建索引以提升查询性能
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_work_orders_trace_process ON work_orders(trace_id, process_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_claimed_orders_trace_process ON claimed_orders(trace_id, process_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_work_results_trace_process ON work_results(trace_id, process_id)')
        
        conn.commit()
        conn.close()
    
    def check_time_limit(self, days: int) -> bool:
        """检查是否超过指定天数限制"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            cursor.execute("SELECT value FROM system_config WHERE key = 'first_start_time'")
            row = cursor.fetchone()
            conn.close()
            
            if row:
                start_time = datetime.strptime(row[0], "%Y-%m-%d %H:%M:%S")
                return datetime.now() > start_time + timedelta(days=days)
            return False
        except:
            return False
    
    def create_work_order(self, order_data: Dict) -> bool:
        """
        创建工单
        :param order_data: 工单数据字典
        :return: 是否成功
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 处理工单中的目标扭矩：原始值为0则删除，否则取整
            bolts = order_data.get('bolts', [])
            processed_bolts = []
            for bolt in bolts:
                processed_bolt = bolt.copy()
                if 'target_torque' in processed_bolt:
                    if processed_bolt['target_torque'] == 0:
                        continue  # 跳过目标扭矩为0的记录
                    processed_bolt['target_torque'] = int(processed_bolt['target_torque'])
                processed_bolts.append(processed_bolt)
            
            # 将bolts列表转换为JSON字符串
            bolts_json = json.dumps(processed_bolts, ensure_ascii=False)
            
            cursor.execute('''
                INSERT OR REPLACE INTO work_orders 
                (trace_id, process_id, process_name, product_name, operator, status, bolts, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                order_data.get('trace_id'),
                order_data.get('process_id'),
                order_data.get('process_name', ''),
                order_data.get('product_name', ''),
                order_data.get('operator', ''),
                order_data.get('status', 'pending'),
                bolts_json,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"创建工单失败: {e}")
            return False
    
    def get_work_order(self, trace_id: str, process_id: str) -> Optional[Dict]:
        """
        获取工单详情
        :param trace_id: 追溯号
        :param process_id: 工序号
        :return: 工单数据字典或None
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM work_orders 
                WHERE trace_id = ? AND process_id = ?
            ''', (trace_id, process_id))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                result = {
                    "trace_id": row['trace_id'],
                    "process_id": row['process_id'],
                    "process_name": row['process_name'],
                    "product_name": row['product_name'],
                    "operator": row['operator'],
                    "status": row['status'],
                    "bolts": json.loads(row['bolts'])
                }
                # 兼容旧数据（如果有station字段）
                if 'station' in row.keys():
                    result['station'] = row['station']
                return result
            return None
        except Exception as e:
            print(f"查询工单失败: {e}")
            return None
    
    def get_available_work_orders(self, trace_id: str, process_id: str) -> List[Dict]:
        """
        获取可用工单列表（未认领的）
        :param trace_id: 追溯号
        :param process_id: 工序号
        :return: 工单列表
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 一次查询获取未认领的工单
            cursor.execute('''
                SELECT wo.* FROM work_orders wo
                LEFT JOIN claimed_orders co ON wo.trace_id = co.trace_id AND wo.process_id = co.process_id
                WHERE wo.trace_id = ? AND wo.process_id = ? AND co.id IS NULL
            ''', (trace_id, process_id))
            
            rows = cursor.fetchall()
            conn.close()
            
            orders = []
            for row in rows:
                order = dict(row)
                order['bolts'] = json.loads(order['bolts'])
                orders.append(order)
            return orders
        except Exception as e:
            print(f"[ERROR] 查询可用工单失败: {e}")
            return []
    
    def get_work_orders(self) -> List[Dict]:
        """
        获取所有工单
        :return: 工单列表
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('SELECT * FROM work_orders')
            rows = cursor.fetchall()
            conn.close()
            
            orders = []
            for row in rows:
                order = {
                    "trace_id": row['trace_id'],
                    "process_id": row['process_id'],
                    "process_name": row['process_name'],
                    "product_name": row['product_name'],
                    "operator": row['operator'],
                    "status": row['status'],
                    "bolts": json.loads(row['bolts'])
                }
                # 兼容旧数据
                if 'station' in row.keys():
                    order['station'] = row['station']
                orders.append(order)
            
            return orders
        except Exception as e:
            print(f"查询所有工单失败: {e}")
            return []
    
    def get_all_available_work_orders(self) -> List[Dict]:
        """
        获取所有可用工单列表（未认领的）
        :return: 工单列表
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 查询所有未认领的工单
            cursor.execute('''
                SELECT wo.* FROM work_orders wo
                LEFT JOIN claimed_orders co ON wo.trace_id = co.trace_id AND wo.process_id = co.process_id
                WHERE co.id IS NULL
                ORDER BY wo.created_at DESC
            ''')
            
            rows = cursor.fetchall()
            conn.close()
            
            orders = []
            for row in rows:
                order = dict(row)
                order['bolts'] = json.loads(order['bolts'])
                orders.append(order)
            return orders
        except Exception as e:
            print(f"[ERROR] 查询所有可用工单失败: {e}")
            import traceback
            traceback.print_exc()
            return []
    
    def get_claimed_order(self, trace_id: str, process_id: str) -> Optional[Dict]:
        """
        获取已认领的工单
        :param trace_id: 追溯号
        :param process_id: 工序号
        :return: 认领信息或None
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM claimed_orders 
                WHERE trace_id = ? AND process_id = ?
            ''', (trace_id, process_id))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                return {
                    "trace_id": row['trace_id'],
                    "process_id": row['process_id'],
                    "process_name": row['process_name'],
                    "operator": row['operator'],
                    "status": row['status'],
                    "claimed_at": row['claimed_at']
                }
            return None
        except Exception as e:
            print(f"查询认领工单失败: {e}")
            return None
    
    def get_all_claimed_orders(self) -> List[Dict]:
        """获取所有已认领的工单"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM claimed_orders WHERE status = "claimed" ORDER BY claimed_at DESC')
            rows = cursor.fetchall()
            conn.close()
            return [dict(row) for row in rows]
        except Exception as e:
            print(f"获取已认领工单失败: {e}")
            return []
    
    def claim_work_order(self, trace_id: str, process_id: str, operator: str) -> bool:
        """
        认领工单
        :param trace_id: 追溯号
        :param process_id: 工序号
        :param operator: 操作员
        :return: 是否成功
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 获取工单信息以获取process_name
            order = self.get_work_order(trace_id, process_id)
            process_name = order.get('process_name', '') if order else ''
            
            cursor.execute('''
                INSERT OR REPLACE INTO claimed_orders 
                (trace_id, process_id, process_name, operator, status, claimed_at)
                VALUES (?, ?, ?, ?, 'claimed', ?)
            ''', (
                trace_id,
                process_id,
                process_name,
                operator,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"认领工单失败: {e}")
            return False
    
    def release_work_order(self, trace_id: str, process_id: str) -> bool:
        """
        释放工单（取消认领）
        :param trace_id: 追溯号
        :param process_id: 工序号
        :return: 是否成功
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                DELETE FROM claimed_orders 
                WHERE trace_id = ? AND process_id = ?
            ''', (trace_id, process_id))
            
            conn.commit()
            deleted = cursor.rowcount > 0
            conn.close()
            return deleted
        except Exception as e:
            print(f"释放工单失败: {e}")
            return False
    
    def unclaim_work_order(self, trace_id: str, process_id: str) -> bool:
        """退领工单"""
        return self.release_work_order(trace_id, process_id)
    
    def delete_work_order(self, trace_id: str, process_id: str) -> bool:
        """删除工单"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            cursor.execute('DELETE FROM work_orders WHERE trace_id = ? AND process_id = ?', (trace_id, process_id))
            cursor.execute('DELETE FROM claimed_orders WHERE trace_id = ? AND process_id = ?', (trace_id, process_id))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"删除工单失败: {e}")
            return False
    
    def submit_work_order(self, trace_id: str, process_id: str, bolts: List[Dict],
                          device_sn: str = None, device_name: str = None) -> bool:
        """
        提交工单结果
        :param trace_id: 追溯号
        :param process_id: 工序号
        :param bolts: 螺栓数据列表
        :param device_sn: 设备SN号
        :param device_name: 设备名称
        :return: 是否成功
        """
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 获取工单信息
            order = self.get_work_order(trace_id, process_id)
            process_name = order.get('process_name', '') if order else ''
            
            # 更新认领状态为已完成
            cursor.execute('''
                UPDATE claimed_orders 
                SET status = 'completed', completed_at = ?
                WHERE trace_id = ? AND process_id = ?
            ''', (
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                trace_id,
                process_id
            ))
            
            # 保存结果
            bolts_json = json.dumps(bolts, ensure_ascii=False)
            cursor.execute('''
                INSERT INTO work_results
                (trace_id, process_id, process_name, bolts, device_sn, device_name, submitted_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                trace_id,
                process_id,
                process_name,
                bolts_json,
                device_sn,
                device_name,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"提交工单失败: {e}")
            return False
    
    def get_work_result(self, trace_id: str, process_id: str) -> Optional[Dict]:
        """
        查询工单执行结果
        :param trace_id: 追溯号
        :param process_id: 工序号
        :return: 执行结果字典，如果不存在则返回None
        """
        try:
            conn = self.get_connection()
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM work_results 
                WHERE trace_id = ? AND process_id = ?
                ORDER BY submitted_at DESC
                LIMIT 1
            ''', (trace_id, process_id))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                # 解析螺栓数据
                try:
                    bolts = json.loads(row['bolts'])
                except:
                    bolts = []
                
                return {
                    "id": row['id'],
                    "trace_id": row['trace_id'],
                    "process_id": row['process_id'],
                    "process_name": row['process_name'],
                    "bolts": bolts,
                    "device_sn": row['device_sn'],
                    "device_name": row['device_name'],
                    "submitted_at": row['submitted_at']
                }
            return None
        except Exception as e:
            print(f"查询工单结果失败: {e}")
            return None
    
    def get_work_results(self, trace_id: str = None, process_id: str = None) -> List[Dict]:
        """
        查询工单执行结果列表
        :param trace_id: 追溯号（可选）
        :param process_id: 工序号（可选）
        :return: 执行结果列表
        """
        try:
            conn = self.get_connection()
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            if trace_id and process_id:
                # 查询指定追溯号和工序号的结果
                cursor.execute('''
                    SELECT * FROM work_results 
                    WHERE trace_id = ? AND process_id = ?
                    ORDER BY submitted_at DESC
                ''', (trace_id, process_id))
            elif trace_id:
                # 只按追溯号查询
                cursor.execute('''
                    SELECT * FROM work_results 
                    WHERE trace_id = ?
                    ORDER BY submitted_at DESC
                ''', (trace_id,))
            elif process_id:
                # 只按工序号查询
                cursor.execute('''
                    SELECT * FROM work_results 
                    WHERE process_id = ?
                    ORDER BY submitted_at DESC
                ''', (process_id,))
            else:
                # 查询所有结果
                cursor.execute('''
                    SELECT * FROM work_results 
                    ORDER BY submitted_at DESC
                ''')
            
            rows = cursor.fetchall()
            conn.close()
            
            results = []
            for row in rows:
                try:
                    bolts = json.loads(row['bolts'])
                except:
                    bolts = []
                
                results.append({
                    "id": row['id'],
                    "trace_id": row['trace_id'],
                    "process_id": row['process_id'],
                    "process_name": row['process_name'],
                    "bolts": bolts,
                    "device_sn": row['device_sn'],
                    "device_name": row['device_name'],
                    "submitted_at": row['submitted_at']
                })
            
            return results
        except Exception as e:
            print(f"查询工单结果列表失败: {e}")
            return []
    
    # ========== 扳手设备管理方法 ==========
    
    def create_wrench_device(self, device_data: Dict) -> bool:
        """创建扳手设备"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT OR REPLACE INTO wrench_devices 
                (device_name, device_sn, ip_address, port, address_code, status, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                device_data.get('device_name'),
                device_data.get('device_sn'),
                device_data.get('ip_address'),
                device_data.get('port', 7888),
                device_data.get('address_code', 1),
                device_data.get('status', 'offline'),
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"创建扳手设备失败: {e}")
            return False
    
    def get_wrench_devices(self) -> List[Dict]:
        """获取所有扳手设备"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('SELECT * FROM wrench_devices ORDER BY device_name')
            rows = cursor.fetchall()
            conn.close()
            
            devices = []
            for row in rows:
                devices.append({
                    "id": row['id'],
                    "device_name": row['device_name'],
                    "device_sn": row['device_sn'],
                    "ip_address": row['ip_address'],
                    "port": row['port'],
                    "address_code": row['address_code'],
                    "status": row['status'],
                    "last_heartbeat": row['last_heartbeat']
                })
            
            return devices
        except Exception as e:
            print(f"查询扳手设备失败: {e}")
            return []
    
    def get_wrench_device(self, device_id: int) -> Optional[Dict]:
        """获取扳手设备详情"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('SELECT * FROM wrench_devices WHERE id = ?', (device_id,))
            row = cursor.fetchone()
            conn.close()
            
            if row:
                return {
                    "id": row['id'],
                    "device_name": row['device_name'],
                    "device_sn": row['device_sn'],
                    "ip_address": row['ip_address'],
                    "port": row['port'],
                    "address_code": row['address_code'],
                    "status": row['status'],
                    "last_heartbeat": row['last_heartbeat']
                }
            return None
        except Exception as e:
            print(f"查询扳手设备失败: {e}")
            return None
    
    def update_wrench_device_status(self, device_id: int, status: str) -> bool:
        """更新扳手设备状态"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                UPDATE wrench_devices 
                SET status = ?, last_heartbeat = ?, updated_at = ?
                WHERE id = ?
            ''', (
                status,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                device_id
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"更新扳手设备状态失败: {e}")
            return False
    
    def update_wrench_device_sn(self, device_id: int, device_sn: str) -> bool:
        """更新扳手设备SN码"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                UPDATE wrench_devices 
                SET device_sn = ?, updated_at = ?
                WHERE id = ?
            ''', (
                device_sn,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                device_id
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"更新扳手设备SN码失败: {e}")
            return False
    
    def update_wrench_device(self, device_id: int, device_data: Dict) -> bool:
        """更新扳手设备信息"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('''
                UPDATE wrench_devices 
                SET device_name = ?, device_sn = ?, ip_address = ?, port = ?, 
                    address_code = ?, updated_at = ?
                WHERE id = ?
            ''', (
                device_data.get('device_name'),
                device_data.get('device_sn'),
                device_data.get('ip_address'),
                device_data.get('port', 7888),
                device_data.get('address_code', 1),
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                device_id
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"更新扳手设备失败: {e}")
            return False
    
    def delete_wrench_device(self, device_id: int) -> bool:
        """删除扳手设备"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM wrench_devices WHERE id = ?', (device_id,))
            
            conn.commit()
            deleted = cursor.rowcount > 0
            conn.close()
            return deleted
        except Exception as e:
            print(f"删除扳手设备失败: {e}")
            return False