"""
title: Usage Tracker
author: InfamyStudio CortexNetwork
author_url: https://github.com/open-webui
funding_url: https://github.com/open-webui
version: 0.1
"""

from pydantic import BaseModel, Field
from typing import Optional
import tiktoken
import json
import os

MODEL_PRICING = {
    "gpt-4o": {
        "input": 5,
        "output": 15,
        "input-tokens": 1000000,
        "output-tokens": 1000000,
    },
    "gpt-4o-mini": {
        "input": 0.150,
        "output": 0.600,
        "input-tokens": 1000000,
        "output-tokens": 1000000,
    },
    "gpt-4-turbo": {
        "input": 10,
        "output": 30,
        "input-tokens": 1000000,
        "output-tokens": 1000000,
    },
    "gpt-4": {
        "input": 30,
        "output": 60,
        "input-tokens": 1000000,
        "output-tokens": 1000000,
    },
    "gpt-3.5-turbo": {
        "input": 3,
        "output": 6,
        "input-tokens": 1000000,
        "output-tokens": 1000000,
    },
    "gpt-3.5-turbo-16k": {
        "input": 3,
        "output": 4,
        "input-tokens": 1000000,
        "output-tokens": 1000000,
    },
    "tts-1": {
        "output": 15,
        "output-tokens": 1000000,
    },
}

class Filter:
    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for the filter operations."
        )
        pass

    def __init__(self):
        self.valves = self.Valves()
        self.input_cost = 0
        self.output_cost = 0
        self.cost_file = 'user_costing.json'
        self._ensure_cost_file_exists()
        pass

    def _ensure_cost_file_exists(self):
        if not os.path.exists(self.cost_file):
            with open(self.cost_file, 'w') as file:
                json.dump({}, file)

    def _read_costs(self):
        with open(self.cost_file, 'r') as file:
            return json.load(file)

    def _write_costs(self, costs):
        with open(self.cost_file, 'w') as file:
            json.dump(costs, file)

    def _update_user_cost(self, user_email, additional_cost):
        costs = self._read_costs()
        if user_email in costs:
            costs[user_email] += additional_cost
        else:
            costs[user_email] = additional_cost
        self._write_costs(costs)

    def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        if __user__.get("role", "admin") in ["user", "admin"]:
            messages = body.get("messages", [])
            model = body.get("model")
            try:
                enc = tiktoken.encoding_for_model(model)
            except Exception as e:
                if "gpt-3.5-turbo" in model:
                    enc = tiktoken.encoding_for_model("gpt-3.5-turbo-")
                else:
                    print(f"Error: {e}")
                    enc = tiktoken.encoding_for_model("gpt-4o")

            if model in MODEL_PRICING:
                total_tokens = sum(
                    len(enc.encode(message["content"])) for message in messages
                )
                model_data = MODEL_PRICING[model]
                self.input_cost = (
                    total_tokens / model_data.get("input-tokens", 1)
                ) * model_data.get("input", 0)

        return body

    def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        messages = body.get("messages", [])
        model = body.get("model")
        try:
            enc = tiktoken.encoding_for_model(model)
        except Exception as e:
            if "gpt-3.5-turbo" in model:
                enc = tiktoken.encoding_for_model("gpt-3.5-turbo-")
            else:
                print(f"Error: {e}")
                enc = tiktoken.encoding_for_model("gpt-4o")

        if model in MODEL_PRICING:
            total_tokens = sum(
                len(enc.encode(message["content"])) for message in messages
            )
            model_data = MODEL_PRICING[model]
            self.output_cost = (
                total_tokens / model_data.get("output-tokens", 1)
            ) * model_data.get("output", 0)
            total_cost = self.input_cost + self.output_cost

            if __user__:
                print(f"User: {__user__}")
                user_email = __user__.get('email')
                self._update_user_cost(user_email, total_cost)

        return body