Backend/db/repositories/robot_repository.py

from typing import List, Optional
from loguru import logger
from db.connection import get_connection
from model.robot import Robot

class RobotRepository:
    def get_all(self) -> List[Robot]:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("SELECT * FROM robots ORDER BY id")
                    robots = [
                        Robot(
                            id=row[0],
                            status=row[1],
                            battery_level=row[2],
                            last_update=row[3],
                            current_zone=row[4],
                            current_row=row[5],
                            current_shelf=row[6]
                        ) for row in cur.fetchall()
                    ]
                    logger.debug(f"Получено {len(robots)} роботов")
                    return robots
        except Exception as e:
            logger.error(f"Ошибка получения всех роботов: {e}")
            return []

    def get_by_id(self, robot_id: str) -> Optional[Robot]:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("SELECT * FROM robots WHERE id = %s", (robot_id,))
                    row = cur.fetchone()
                    if row:
                        logger.debug(f"Робот {robot_id} найден")
                        return Robot(*row)
                    logger.warning(f"Робот {robot_id} не найден")
                    return None
        except Exception as e:
            logger.error(f"Ошибка получения робота {robot_id}: {e}")
            return None

    def update_robot(
        self,
        robot_id: str,
        status: str = None,
        battery_level: int = None,
        current_zone: str = None,
        current_row: int = None,
        current_shelf: int = None
    ) -> bool:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    updates = ["last_update = CURRENT_TIMESTAMP"]
                    params = []

                    if status is not None:
                        updates.append("status = %s")
                        params.append(status)

                    if battery_level is not None:
                        updates.append("battery_level = %s")
                        params.append(battery_level)

                    if current_zone is not None:
                        updates.append("current_zone = %s")
                        params.append(current_zone)

                    if current_row is not None:
                        updates.append("current_row = %s")
                        params.append(current_row)

                    if current_shelf is not None:
                        updates.append("current_shelf = %s")
                        params.append(current_shelf)

                    params.append(robot_id)
                    query = f"UPDATE robots SET {', '.join(updates)} WHERE id = %s"

                    cur.execute(query, params)
                    conn.commit()
                    success = cur.rowcount > 0
                    if success:
                        logger.debug(f"Робот {robot_id} успешно обновлен")
                    else:
                        logger.warning(f"Робот {robot_id} не найден для обновления")
                    return success
        except Exception as e:
            logger.error(f"Ошибка обновления робота {robot_id}: {e}")
            return False

    def get_robots_by_status(self, status: str) -> List[Robot]:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("SELECT * FROM robots WHERE status = %s ORDER BY id", (status,))
                    robots = [Robot(*row) for row in cur.fetchall()]
                    logger.debug(f"Получено {len(robots)} роботов со статусом {status}")
                    return robots
        except Exception as e:
            logger.error(f"Ошибка получения роботов по статусу {status}: {e}")
            return []

    def get_low_battery_robots(self, threshold: int = 20) -> List[Robot]:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("""
                        SELECT * FROM robots 
                        WHERE battery_level < %s AND status = 'active'
                        ORDER BY battery_level ASC
                    """, (threshold,))
                    robots = [Robot(*row) for row in cur.fetchall()]
                    logger.debug(f"Получено {len(robots)} роботов с низким зарядом (<{threshold}%)")
                    return robots
        except Exception as e:
            logger.error(f"Ошибка получения роботов с низким зарядом: {e}")
            return []

    def get_robots_in_zone(self, zone: str) -> List[Robot]:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("SELECT * FROM robots WHERE current_zone = %s ORDER BY id", (zone,))
                    robots = [Robot(*row) for row in cur.fetchall()]
                    logger.debug(f"Получено {len(robots)} роботов в зоне {zone}")
                    return robots
        except Exception as e:
            logger.error(f"Ошибка получения роботов в зоне {zone}: {e}")
            return []

    def create_robot(
        self,
        robot_id: str,
        status: str = 'active',
        battery_level: int = 100
    ) -> bool:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("""
                        INSERT INTO robots (id, status, battery_level, last_update)
                        VALUES (%s, %s, %s, CURRENT_TIMESTAMP)
                    """, (robot_id, status, battery_level))
                    conn.commit()
                    success = cur.rowcount > 0
                    if success:
                        logger.debug(f"Создан новый робот {robot_id}")
                    else:
                        logger.warning(f"Не удалось создать робота {robot_id}")
                    return success
        except Exception as e:
            logger.error(f"Ошибка создания робота {robot_id}: {e}")
            return False

    def delete_robot(self, robot_id: str) -> bool:
        try:
            with get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("DELETE FROM robots WHERE id = %s", (robot_id,))
                    conn.commit()
                    success = cur.rowcount > 0
                    if success:
                        logger.debug(f"Робот {robot_id} удален")
                    else:
                        logger.warning(f"Робот {robot_id} не найден для удаления")
                    return success
        except Exception as e:
            logger.error(f"Ошибка удаления робота {robot_id}: {e}")
            return False