from connection.connection import Database

class WAFModel:
    def __init__(self):
        self.db = Database()

    def create_waf_metric(self, web_acl_id, waf_allowed_requests, waf_blocked_requests, start_time, end_time):
        query = """
        INSERT INTO waf_metrics (web_acl_id, waf_allowed_requests, waf_blocked_requests, start_time, end_time) 
        VALUES (%s, %s, %s, %s, %s) RETURNING id;
        """
        try:
            return self.db.execute(query, (web_acl_id, waf_allowed_requests, waf_blocked_requests, start_time, end_time), fetchone=True)
        except Exception as e:
            print(f"Error creating WAF metric: {e}")
            return None

    def get_waf_metric(self, metric_id):
        query = "SELECT id, web_acl_id, waf_allowed_requests, waf_blocked_requests, start_time, end_time FROM waf_metrics WHERE id = %s;"
        try:
            return self.db.execute(query, (metric_id,), fetchone=True)
        except Exception as e:
            print(f"Error fetching WAF metric: {e}")
            return None

    def update_waf_metric(self, metric_id, waf_allowed_requests=None, waf_blocked_requests=None):
        query = """
        UPDATE waf_metrics SET 
            waf_allowed_requests = COALESCE(%s, waf_allowed_requests),
            waf_blocked_requests = COALESCE(%s, waf_blocked_requests)
        WHERE id = %s RETURNING id;
        """
        try:
            return self.db.execute(query, (waf_allowed_requests, waf_blocked_requests, metric_id), fetchone=True)
        except Exception as e:
            print(f"Error updating WAF metric: {e}")
            return None

    def delete_waf_metric(self, metric_id):
        query = "DELETE FROM waf_metrics WHERE id = %s RETURNING id;"
        try:
            return self.db.execute(query, (metric_id,), fetchone=True)
        except Exception as e:
            print(f"Error deleting WAF metric: {e}")
            return None