from .BaseLLM import BaseLLM
import os
from openai import OpenAI
import tiktoken
encoding = tiktoken.encoding_for_model("gpt-4o")

class OpenRouter(BaseLLM):
    def __init__(self, model="deepseek/deepseek-r1:free"):
        super(OpenRouter, self).__init__()
        self.client = OpenAI(
        api_key=os.getenv("OPENROUTER_API_KEY"), 
        base_url="https://openrouter.ai/api/v1",
        )
        self.model_name = model
        self.messages = []
        self.in_token = 0
        self.out_token = 0
    def initialize_message(self):
        self.messages = []

    def ai_message(self, payload):
        self.messages.append({"role": "ai", "content": payload})

    def system_message(self, payload):
        self.messages.append({"role": "system", "content": payload})

    def user_message(self, payload):
        self.messages.append({"role": "user", "content": payload})

    def get_response(self,temperature = 0.8):
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=self.messages
            )
        return completion.choices[0].message.content
    
    def chat(self,text,temperature = 0.8):
        self.initialize_message()
        self.user_message(text)
        response = self.get_response(temperature = temperature)
        print("In",self.count_token(text))
        print("Out", self.count_token(response))
        self.in_token += self.count_token(text)
        self.out_token += self.count_token(response)
        return response
    
    def print_prompt(self):
        for message in self.messages:
            print(message)
    def count_token(self,text,):
        return len(encoding.encode(text))