# -*- coding: utf-8 -*-
# File: database_handler.py
# Purpose: 处理与数据库的交互，包括连接、数据读写（如用户信息、无人机数据、任务记录等）。

import mysql.connector
from mysql.connector import Error
from typing import Dict, List, Any, Optional
import json
from datetime import datetime

class DatabaseHandler:
    #替换数据库账号密码
    def __init__(self, host: str, user: str, password: str, database: str):
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.connection = None
        self.cursor = None

    def connect(self) -> bool:
        """连接到MySQL数据库"""
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database
            )
            self.cursor = self.connection.cursor(dictionary=True)
            print(f"数据库连接成功: {self.database}")
            return True
        except Error as e:
            print(f"数据库连接失败: {str(e)}")
            return False

    def disconnect(self):
        """断开数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()

    def execute_query(self, query: str, params: Optional[tuple] = None) -> List[Dict]:
        """执行查询并返回结果"""
        try:
            self.cursor.execute(query, params)
            return self.cursor.fetchall()
        except Error as e:
            print(f"查询执行失败: {str(e)}")
            return []

    def execute_update(self, query: str, params: Optional[tuple] = None) -> bool:
        """执行更新操作"""
        try:
            self.cursor.execute(query, params)
            self.connection.commit()
            return True
        except Error as e:
            print(f"更新执行失败: {str(e)}")
            self.connection.rollback()
            return False

    def save_mavlink_message(self, msg_type: str, msg_data: Dict) -> bool:
        """保存MAVLink消息到数据库"""
        query = """
        INSERT INTO mavlink_messages 
        (message_type, message_data, timestamp) 
        VALUES (%s, %s, %s)
        """
        params = (
            msg_type,
            json.dumps(msg_data),
            datetime.now()
        )
        return self.execute_update(query, params)

    def get_mavlink_messages(self, msg_type: Optional[str] = None, 
                           start_time: Optional[datetime] = None,
                           end_time: Optional[datetime] = None) -> List[Dict]:
        """获取MAVLink消息"""
        query = "SELECT * FROM mavlink_messages WHERE 1=1"
        params = []
        
        if msg_type:
            query += " AND message_type = %s"
            params.append(msg_type)
        if start_time:
            query += " AND timestamp >= %s"
            params.append(start_time)
        if end_time:
            query += " AND timestamp <= %s"
            params.append(end_time)
            
        query += " ORDER BY timestamp DESC"
        return self.execute_query(query, tuple(params) if params else None)

    def create_tables(self):
        """创建必要的数据库表"""
        queries = [
            """
            CREATE TABLE IF NOT EXISTS mavlink_messages (
                id INT AUTO_INCREMENT PRIMARY KEY,
                message_type VARCHAR(50) NOT NULL,
                message_data JSON NOT NULL,
                timestamp DATETIME NOT NULL,
                INDEX idx_message_type (message_type),
                INDEX idx_timestamp (timestamp)
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS drone_status (
                id INT AUTO_INCREMENT PRIMARY KEY,
                drone_id VARCHAR(50) NOT NULL,
                latitude FLOAT,
                longitude FLOAT,
                altitude FLOAT,
                battery_level FLOAT,
                status VARCHAR(50),
                timestamp DATETIME NOT NULL,
                INDEX idx_drone_id (drone_id),
                INDEX idx_timestamp (timestamp)
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS drone_connections (
                id INT AUTO_INCREMENT PRIMARY KEY,
                drone_id VARCHAR(50) NOT NULL UNIQUE,
                ip_address VARCHAR(50) NOT NULL,
                port INT NOT NULL,
                connection_type VARCHAR(20) NOT NULL,
                is_active BOOLEAN DEFAULT TRUE,
                last_connected DATETIME,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                INDEX idx_drone_id (drone_id),
                INDEX idx_is_active (is_active)
            )
            """
        ]
        
        for query in queries:
            try:
                self.cursor.execute(query)
                self.connection.commit()
            except Error as e:
                print(f"创建表失败: {str(e)}")
                self.connection.rollback()

    def get_all_active_drones(self) -> List[Dict[str, Any]]:
        """获取所有活跃的无人机连接信息"""
        query = """
        SELECT * FROM drone_connections 
        WHERE is_active = TRUE
        """
        return self.execute_query(query)

    def get_drone_connection(self, drone_id: str) -> Optional[Dict[str, Any]]:
        """获取指定无人机的连接信息"""
        query = """
        SELECT * FROM drone_connections 
        WHERE drone_id = %s AND is_active = TRUE
        """
        result = self.execute_query(query, (drone_id,))
        return result[0] if result else None

    def update_drone_connection_status(self, drone_id: str, is_active: bool) -> bool:
        """更新无人机连接状态"""
        query = """
        UPDATE drone_connections 
        SET is_active = %s, last_connected = %s 
        WHERE drone_id = %s
        """
        params = (is_active, datetime.now(), drone_id)
        return self.execute_update(query, params)

    def add_drone_connection(self, drone_id: str, ip_address: str, port: int, 
                           connection_type: str = 'udp') -> bool:
        """添加新的无人机连接信息"""
        query = """
        INSERT INTO drone_connections 
        (drone_id, ip_address, port, connection_type) 
        VALUES (%s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE 
        ip_address = VALUES(ip_address),
        port = VALUES(port),
        connection_type = VALUES(connection_type),
        is_active = TRUE,
        updated_at = CURRENT_TIMESTAMP
        """
        params = (drone_id, ip_address, port, connection_type)
        return self.execute_update(query, params)