"""
title: Simple PII Redaction Filter
author: justinh-rahb
author_url: https://github.com/justinh-rahb
funding_url: https://github.com/open-webui
version: 0.1.0-rc1
license: MIT
"""

import re
from pydantic import BaseModel, Field
from typing import Optional


class Filter:
    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for the filter operations."
        )
        redact_email: bool = Field(default=True, description="Redact email addresses.")
        redact_phone: bool = Field(default=True, description="Redact phone numbers.")
        redact_ssn: bool = Field(
            default=True, description="Redact social security numbers."
        )
        redact_credit_card: bool = Field(
            default=True, description="Redact credit card numbers."
        )
        redact_ip_address: bool = Field(
            default=True, description="Redact IP addresses."
        )
        enabled_for_admins: bool = Field(
            default=False,
            description="Whether PII Redaction is enabled for admin users.",
        )

    def __init__(self):
        self.file_handler = False
        self.valves = self.Valves()
        self.patterns = {
            "email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"),
            "phone": re.compile(r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b"),
            "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
            "credit_card": re.compile(r"\b(?:\d{4}[-\s]?){3}\d{4}\b"),
            "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
        }

    def redact_pii(self, text: str) -> str:
        if self.valves.redact_email:
            text = self.patterns["email"].sub("[EMAIL REDACTED]", text)
        if self.valves.redact_phone:
            text = self.patterns["phone"].sub("[PHONE REDACTED]", text)
        if self.valves.redact_ssn:
            text = self.patterns["ssn"].sub("[SSN REDACTED]", text)
        if self.valves.redact_credit_card:
            text = self.patterns["credit_card"].sub("[CREDIT CARD REDACTED]", text)
        if self.valves.redact_ip_address:
            text = self.patterns["ip_address"].sub("[IP ADDRESS REDACTED]", text)
        return text

    def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        print(f"inlet:{__name__}")
        print(f"inlet:body:{body}")
        print(f"inlet:user:{__user__}")

        if (
            __user__ is None
            or not __user__.get("role") == "admin"
            or self.valves.enabled_for_admins
        ):
            messages = body.get("messages", [])
            for message in messages:
                if message.get("role") == "user":
                    message["content"] = self.redact_pii(message["content"])

        return body

    def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        print(f"outlet:{__name__}")
        print(f"outlet:body:{body}")
        print(f"outlet:user:{__user__}")

        return body