import requests
import os
import json

SYSTEM_PROMPT = """
Tôi là AI STV, một trợ lý AI thông minh, thân thiện, được phát triển bởi Trọng Phúc.
Luôn giao tiếp bằng tiếng Việt hoặc ngôn ngữ khác, câu trả lời ngắn gọn, dễ hiểu và hữu ích.
"""

# Danh sách token test dùng cho người mới
TEST_TOKENS = [
    "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"
]

# Cấu hình
USAGE_LOG = "ip_usage.json"
MAX_FREE_QUESTIONS = 5

def get_user_ip():
    try:
        response = requests.get("https://api.ipify.org?format=json", timeout=5)
        return response.json()["ip"]
    except:
        return "unknown"

def load_usage():
    if os.path.exists(USAGE_LOG):
        with open(USAGE_LOG, "r") as f:
            return json.load(f)
    return {}

def save_usage(data):
    with open(USAGE_LOG, "w") as f:
        json.dump(data, f)

class STVBot:
    def __init__(self, user_token=None, system_prompt=SYSTEM_PROMPT):
        self.api_url = "https://openrouter.ai/api/v1/chat/completions"
        self.system_prompt = system_prompt
        self.model = "nousresearch/deephermes-3-mistral-24b-preview:free"
        self.history = []
        if system_prompt:
            self.history.append({"role": "system", "content": system_prompt})

        self.user_ip = get_user_ip()
        self.token = user_token.strip() if user_token and user_token.strip() != "" else None
        self.free_mode = self.token is None
        self.usage_data = load_usage()

        if self.free_mode:
            # nếu chưa có IP trong usage_data thì khởi tạo 0
            if self.user_ip not in self.usage_data:
                self.usage_data[self.user_ip] = 0

    def get_valid_test_token(self):
        for token in TEST_TOKENS:
            if self.validate_token(token):
                return token
        return TEST_TOKENS[0]

    def validate_token(self, token):
        # Có thể thêm kiểm tra token ở đây
        return True

    def ask(self, prompt):
        # Nếu dùng free mode thì kiểm tra lượt hỏi
        if self.free_mode:
            count = self.usage_data.get(self.user_ip, 0)
            if count >= MAX_FREE_QUESTIONS:
                return (
                    "⚠️ Bạn đã sử dụng hết 5 lượt hỏi miễn phí.\n"
                    "Để tiếp tục sử dụng trợ lý AI STV, vui lòng nhập token riêng của bạn.\n"
                    "Nếu chưa có token, bạn có thể tham gia Discord để được hỗ trợ miễn phí:\n"
                    "👉 https://discord.gg/Ze7RTExgdv"
                )
            # cấp token test
            self.token = self.get_valid_test_token()
            # tăng lượt hỏi
            self.usage_data[self.user_ip] = count + 1
            save_usage(self.usage_data)

        self.history.append({"role": "user", "content": prompt})
        body = {
            "model": self.model,
            "messages": self.history,
            "max_tokens": 300,
            "temperature": 0.7,
        }
        headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://openrouter.ai",
            "X-Title": "STV Chat"
        }

        try:
            response = requests.post(self.api_url, headers=headers, json=body, timeout=15)
            data = response.json()
            reply = data["choices"][0]["message"]["content"]
            self.history.append({"role": "assistant", "content": reply})
            return reply.strip()
        except Exception as e:
            return f"⚠️ Lỗi khi kết nối API: {e}"