# database/repositories/robot_repository.py
from typing import List, Optional
from database.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")
                    return [
                        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()
                    ]
        except Exception as e:
            print(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()
                    return Robot(*row) if row else None
        except Exception as e:
            print(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()
                    return cur.rowcount > 0
        except Exception as e:
            print(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,))
                    return [Robot(*row) for row in cur.fetchall()]
        except Exception as e:
            print(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,))
                    return [Robot(*row) for row in cur.fetchall()]
        except Exception as e:
            print(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,))
                    return [Robot(*row) for row in cur.fetchall()]
        except Exception as e:
            print(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()
                    return cur.rowcount > 0
        except Exception as e:
            print(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()
                    return cur.rowcount > 0
        except Exception as e:
            print(f"Ошибка удаления робота {robot_id}: {e}")
            return False