#!/usr/bin/python
# -*- coding: utf-8 -*-

# --------------------------------------------------------------------------- #
#                                                                             #
#    Plugin for iSida Jabber Bot                                              #
#    AI Assistant with Ollama + Roles + Context                              #
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.    #
#                                                                             #
# --------------------------------------------------------------------------- #

import re
import json
import random
import time
import urllib.parse
import urllib.request

# ============== НАСТРОЙКИ ==============
AI_CONFIG = {
    'enabled': True,
    'respond_to_mention': True,
    'respond_to_questions': True,
    'max_response_length': 512,
    'use_calculator': True,
    'use_weather': True,
    'model': 'qwen2.5:3b',
    'context_size': 7,
    'ollama_url': 'http://localhost:11434/api/chat',
}

conversation_cache = {}
last_message_time = {}
user_cache = {}
owner_jid = None

# ============== ИНИЦИАЛИЗАЦИЯ ==============

def init_owner():
    """Получает JID владельца бота из конфига"""
    global owner_jid
    try:
        owner_jid = SuperAdmin.lower()
    except:
        owner_jid = None

# ============== ТРАНСЛИТЕРАЦИЯ ==============

RUS_TO_LAT = {
    'а': 'a', 'б': 'b', 'в': 'v', 'г': 'g', 'д': 'd', 'е': 'e', 'ё': 'yo',
    'ж': 'zh', 'з': 'z', 'и': 'i', 'й': 'y', 'к': 'k', 'л': 'l', 'м': 'm',
    'н': 'n', 'о': 'o', 'п': 'p', 'р': 'r', 'с': 's', 'т': 't', 'у': 'u',
    'ф': 'f', 'х': 'kh', 'ц': 'ts', 'ч': 'ch', 'ш': 'sh', 'щ': 'shch',
    'ъ': '', 'ы': 'y', 'ь': '', 'э': 'e', 'ю': 'yu', 'я': 'ya',
    'А': 'A', 'Б': 'B', 'В': 'V', 'Г': 'G', 'Д': 'D', 'Е': 'E', 'Ё': 'Yo',
    'Ж': 'Zh', 'З': 'Z', 'И': 'I', 'Й': 'Y', 'К': 'K', 'Л': 'L', 'М': 'M',
    'Н': 'N', 'О': 'O', 'П': 'P', 'Р': 'R', 'С': 'S', 'Т': 'T', 'У': 'U',
    'Ф': 'F', 'Х': 'Kh', 'Ц': 'Ts', 'Ч': 'Ch', 'Ш': 'Sh', 'Щ': 'Shch',
    'Ъ': '', 'Ы': 'Y', 'Ь': '', 'Э': 'E', 'Ю': 'Yu', 'Я': 'Ya'
}

def transliterate(text):
    result = []
    for char in text:
        result.append(RUS_TO_LAT.get(char, char))
    return ''.join(result)

# ============== ПОГОДА ==============

def get_weather(city='Engels', lang='ru'):
    try:
        city_lat = transliterate(city)
        url = f'https://wttr.in/{urllib.parse.quote(city_lat)}?lang={lang}&format=%C+%t+%w+%h'
        req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = resp.read().decode('utf-8').strip()
            if data.startswith('<!DOCTYPE') or data.startswith('<html'):
                return None
            parts = data.split()
            if len(parts) >= 4:
                weather = parts[0]
                temp = parts[1]
                wind = parts[2]
                humid = parts[3]
                return f"{weather}, {temp}, ветер {wind}, влажность {humid}"
            return data
    except:
        return None

def parse_weather_city(text):
    patterns = [
        r'(?:погода|weather).*?(?:в|in)\s+([а-яА-Яa-zA-Z-]+)',
        r'(?:в|in)\s+([а-яА-Яa-zA-Z-]+).*?(?:погода|weather)',
    ]
    for pattern in patterns:
        match = re.search(pattern, text, re.I)
        if match:
            return match.group(1).strip()
    return None

def is_weather_request(text):
    keywords = ['погода', 'weather', 'температура', 'дождь', 'снег', 'ветер']
    return any(kw in text.lower() for kw in keywords)

# ============== РОЛИ ==============

def get_user_info(room, nick):
    """Получает информацию о пользователе с правильным определением ролей"""
    global user_cache, owner_jid

    if nick in user_cache and user_cache[nick].get('time', 0) > time.time() - 300:
        return user_cache[nick]

    try:
        level, jid = get_level(room, nick)
        jid_room = getRoom(jid) if jid else None
        jid_lower = jid.lower() if jid else ''

        # Определяем роль строго по уровням доступа iSida
        if owner_jid and jid_lower == owner_jid:
            # SuperAdmin — владелец бота (уровень 9)
            role = 'superadmin'
            title = random.choice(['Создатель', 'Демиург', 'Архитектор', 'Прародитель'])
        elif level >= 8:
            # Owner комнаты (уровень 8)
            role = 'owner'
            title = random.choice(['Владыка', 'Повелитель', 'Хозяин', 'Князь'])
        elif level >= 7:
            role = 'admin'
            title = random.choice(['Боярин', 'Граф', 'Воевода', 'Князь'])
        elif level >= 5:
            role = 'moderator'
            title = random.choice(['Жандарм', 'Сотник', 'Дружинник', 'Страж'])
        elif level >= 3:
            role = 'member'
            title = random.choice(['Послушник', 'Фримен', 'Свой'])
        elif level >= 0:
            role = 'visitor'
            title = random.choice(['Странник', 'Паломник', 'Гость'])
        else:
            role = 'none'
            title = random.choice(['Немой', 'Безголосый', 'Бесправный'])

        info = {
            'level': level,
            'jid': jid,
            'room': jid_room,
            'role': role,
            'title': title,
            'nick': nick,
            'time': time.time()
        }
        user_cache[nick] = info
        return info
    except:
        return {'level': 0, 'jid': None, 'room': None, 'role': 'unknown', 'title': 'Пользователь', 'nick': nick, 'time': time.time()}

def get_user_title(room, nick):
    info = get_user_info(room, nick)
    return info.get('title', 'Пользователь')

def get_user_role(room, nick):
    info = get_user_info(room, nick)
    return info.get('role', 'unknown')

def get_user_level(room, nick):
    info = get_user_info(room, nick)
    return info.get('level', 0)

# ============== СИСТЕМНЫЙ ПРОМПТ ==============

def get_system_prompt(room, nick):
    """Динамический промпт с информацией о пользователе и контексте"""
    title = get_user_title(room, nick)
    role = get_user_role(room, nick)
    level = get_user_level(room, nick)

    # Формируем описание ролей для контекста
    roles_desc = ""
    for n in user_cache:
        if user_cache[n].get('time', 0) > time.time() - 600:
            u_title = user_cache[n].get('title', 'Пользователь')
            u_role = user_cache[n].get('role', 'unknown')
            roles_desc += f"- {n}: {u_title} ({u_role})\n"

    # Информация о владельце бота
    owner_info = ""
    if owner_jid:
        for n in user_cache:
            if user_cache[n].get('jid', '').lower() == owner_jid:
                owner_info = f"Владелец бота: {n} (Создатель)"
                break

    return f"""Ты — бот-помощник Isida в Jabber-чате.

Информация о пользователе {nick}:
- Роль: {role} (уровень {level})
- Титул: {title}
{owner_info}

Известные пользователи чата:
{roles_desc}

Ты админ чата. Твоя задача — помогать пользователям, отвечать на вопросы, быть дружелюбным и с чувством юмора.

Правила:
1. Отвечай ТОЛЬКО если к тебе обратились по имени или задали прямой вопрос
2. Помни контекст разговора (последние {AI_CONFIG['context_size']} сообщений)
3. Отвечай кратко (2-4 предложения), но по делу
4. Если не знаешь — скажи "Не знаю"
5. Используй титулы пользователей при обращении к ним
6. Будь естественным, не слишком официальным

Ты знаешь роли всех активных пользователей чата и можешь использовать эту информацию в общении."""

# ============== AI ==============

def get_ai_response_ollama(room, jid, nick, messages, max_length):
    try:
        last_msg = messages[-1]['content'] if messages else ''
        system_prompt = get_system_prompt(room, nick)

        # Добавляем контекст из последних сообщений
        context = ""
        if room in conversation_cache and len(conversation_cache[room]) > 1:
            history = conversation_cache[room][-AI_CONFIG['context_size']:]
            for msg in history[:-1]:
                if msg.get('role') == 'user':
                    context += f"{msg.get('nick', 'Пользователь')}: {msg['content']}\n"
                elif msg.get('role') == 'assistant':
                    context += f"Isida: {msg['content']}\n"

        if context:
            full_messages = [
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': f"Контекст разговора:\n{context}\nВопрос: {last_msg}"}
            ]
        else:
            full_messages = [
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': last_msg}
            ]

        data = {
            'model': AI_CONFIG.get('model', 'qwen2.5:3b'),
            'messages': full_messages,
            'stream': False,
            'options': {
                'num_predict': 350,
                'temperature': 0.7,
                'seed': 42,
            }
        }

        ollama_url = AI_CONFIG.get('ollama_url', 'http://localhost:11434/api/chat')
        req = urllib.request.Request(
            ollama_url,
            data=json.dumps(data).encode('utf-8'),
            headers={'Content-Type': 'application/json'}
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            result = json.loads(resp.read().decode('utf-8'))
            answer = result['message']['content']
            if len(answer) > 500:
                answer = answer[:500] + '...'
            return answer
    except Exception as e:
        # Log error for debugging
        try:
            printlog(f'AI Ollama error: {e}')
        except:
            pass
        return None

def get_ai_response(room, jid, nick, messages, max_length=512):
    result = get_ai_response_ollama(room, jid, nick, messages, max_length)
    if result and len(result) > 2:
        return result.strip()
    return random.choice(['Не знаю 😅', 'Я хз, че там по чем', 'Не понял вопроса'])

# ============== ВСПОМОГАТЕЛЬНЫЕ ==============

def is_direct_question_to_bot(text, bot_nick):
    """Проверяет, адресован ли вопрос боту"""
    text_lower = text.lower()
    bot_lower = bot_nick.lower()

    if bot_lower in text_lower:
        return True

    words = text_lower.split()
    if words:
        q_words = ['что', 'как', 'где', 'когда', 'почему', 'зачем', 'сколько', 'кто', 'чей', 'какой', 'какая', 'какое']
        if words[0] in q_words:
            # Это прямой вопрос, начинающийся с вопросительного слова
            return True

    return False

def is_calculation(text):
    patterns = [
        r'\d+\s*[+\-*/]\s*\d+',
        r'\d+\s*(плюс|минус|умножить|разделить|делить)\s*\d+',
        r'сколько\s+будет\s*[\d+\-*/()\s]+',
        r'посчитай\s*[\d+\-*/()\s]+',
    ]
    return any(re.search(p, text, re.I) for p in patterns)

def is_currency(text):
    currencies = ['usd', 'eur', 'rub', 'uah', 'gbp', 'cny', 'jpy', 'chf', 'cad', 'aud']
    return any(c in text.lower() for c in currencies)

def calculate(text):
    try:
        expr = text.lower()
        expr = re.sub(r'сколько\s+будет\s*', '', expr)
        expr = re.sub(r'посчитай\s*', '', expr)
        expr = expr.replace('плюс', '+').replace('минус', '-')
        expr = expr.replace('умножить', '*').replace('разделить', '/')
        expr = expr.replace('делить', '/')
        expr = re.sub(r'\s+', '', expr)
        # Support decimals and more operators
        match = re.search(r'[\d+\-*/().]+', expr)
        if not match:
            return None
        # Safe evaluation without eval - use a simple parser
        expr_clean = match.group().strip()
        # Basic safety check - only allow numbers, operators, parentheses, and dots
        if not re.match(r'^[\d+\-*/().]+$', expr_clean):
            return None
        # Use eval but with restricted globals/locals for safety
        result = eval(expr_clean, {"__builtins__": {}}, {})
        return str(round(float(result), 2))
    except:
        return None

def get_currency_rate(from_cur, to_cur):
    try:
        url = f'https://api.exchangerate-api.com/v4/latest/{from_cur.upper()}'
        req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read().decode('utf-8'))
            return data['rates'].get(to_cur.upper())
    except:
        return None

# ============== ПЕРЕВОД ==============

def handle_translation(room, jid, nick, clean_text, type):
    try:
        text_to_translate = re.sub(r'^(?:переведи|translate)\s+(?:на\s+[а-яa-z]+\s+)?', '', clean_text, flags=re.I)
        text_to_translate = text_to_translate.strip()

        if not text_to_translate:
            return False

        if 'на английский' in clean_text.lower() or 'на english' in clean_text.lower():
            target = 'английский'
        else:
            target = 'русский'

        prompt = f"Переведи на {target} язык: {text_to_translate}"
        messages = [{'role': 'user', 'content': prompt}]
        answer = get_ai_response(room, jid, nick, messages, 300)

        if answer:
            send_msg(type, room, '', answer)
            return True
        return False
    except:
        return False

# ============== ОБРАБОТКА ==============

def ai_process_message(room, jid, nick, type, text):
    if not AI_CONFIG['enabled'] or type != 'groupchat':
        return False

    bot_nick = get_xnick(room)
    if nick == bot_nick:
        return False

    is_mentioned = bot_nick.lower() in text.lower()
    is_direct = is_direct_question_to_bot(text, bot_nick)

    if not is_mentioned and not is_direct:
        return False

    current_time = time.time()
    if room in last_message_time and current_time - last_message_time[room] < 2:
        return False
    last_message_time[room] = current_time

    clean_text = text
    if bot_nick.lower() in clean_text.lower():
        clean_text = re.sub(rf'{bot_nick}[,:\s]*', '', clean_text, flags=re.I).strip()

    if not clean_text:
        return False

    # Обновляем кэш пользователя
    get_user_info(room, nick)

    # Сохраняем в кэш
    if room not in conversation_cache:
        conversation_cache[room] = []
    conversation_cache[room].append({'role': 'user', 'content': clean_text, 'nick': nick})
    if len(conversation_cache[room]) > 20:
        conversation_cache[room] = conversation_cache[room][-20:]

    # ===== ПЕРЕВОД =====
    if re.match(r'^(?:переведи|translate)\s+', clean_text, re.I):
        if handle_translation(room, jid, nick, clean_text, type):
            return True

    # ===== ПОГОДА =====
    if AI_CONFIG['use_weather'] and is_weather_request(clean_text):
        city = parse_weather_city(clean_text)
        if not city:
            city = 'Engels'
        weather = get_weather(city, 'ru')
        if weather:
            send_msg(type, room, '', L('Погода в %s: %s' % (city, weather), '%s/%s' % (room, nick)))
            return True
        send_msg(type, room, '', L('Не могу найти погоду для: %s' % city, '%s/%s' % (room, nick)))
        return True

    # ===== КАЛЬКУЛЯТОР =====
    if AI_CONFIG['use_calculator'] and is_calculation(clean_text):
        result = calculate(clean_text)
        if result:
            send_msg(type, room, '', L('Ответ: %s' % result, '%s/%s' % (room, nick)))
            return True

    # ===== ВАЛЮТЫ =====
    if is_currency(clean_text):
        currencies = re.findall(r'(usd|eur|rub|uah|gbp|cny|jpy|chf|cad|aud)', clean_text, re.I)
        if len(currencies) >= 2:
            from_cur = currencies[0].upper()
            to_cur = currencies[1].upper()
            rate = get_currency_rate(from_cur, to_cur)
            if rate:
                numbers = re.findall(r'(\d+\.?\d*)', clean_text)
                if numbers:
                    amount = float(numbers[0])
                    result = amount * rate
                    send_msg(type, room, '', L('%s %s = %s %s' % (amount, from_cur, round(result, 2), to_cur), '%s/%s' % (room, nick)))
                    return True

    # ===== AI =====
    messages = [{'role': 'user', 'content': clean_text}]
    answer = get_ai_response(room, jid, nick, messages, AI_CONFIG['max_response_length'])
    if answer:
        conversation_cache[room].append({'role': 'assistant', 'content': answer})
        send_msg(type, room, '', answer)
        return True

    return False

def ai_command(type, jid, nick, text):
    if not text:
        send_msg(type, jid, nick, L('Использование: ai <текст>', '%s/%s' % (jid, nick)))
        return

    messages = [{'role': 'user', 'content': text}]
    answer = get_ai_response(jid, jid, nick, messages, 512)
    send_msg(type, jid, nick, answer if answer else L('Не знаю', '%s/%s' % (jid, nick)))

def ai_config(type, jid, nick, text):
    if not text:
        msg = L('Настройки AI:\n', '%s/%s' % (jid, nick))
        for key, value in AI_CONFIG.items():
            msg += '%s: %s\n' % (key, value)
        send_msg(type, jid, nick, msg)
        return

    parts = text.strip().split(' ', 1)
    if len(parts) != 2:
        send_msg(type, jid, nick, L('Использование: aiconfig <ключ> <значение>', '%s/%s' % (jid, nick)))
        return

    key, value = parts[0], parts[1]
    if key not in AI_CONFIG:
        send_msg(type, jid, nick, L('Неизвестный ключ: %s', '%s/%s' % (jid, nick)) % key)
        return

    if value.lower() in ['true', 'on', 'yes', '1']:
        AI_CONFIG[key] = True
    elif value.lower() in ['false', 'off', 'no', '0']:
        AI_CONFIG[key] = False
    else:
        try:
            AI_CONFIG[key] = int(value)
        except:
            AI_CONFIG[key] = value

    send_msg(type, jid, nick, L('Обновлено: %s = %s', '%s/%s' % (jid, nick)) % (key, AI_CONFIG[key]))

# ============== РЕГИСТРАЦИЯ ==============

global execute, message_act_control

init_owner()

execute = [
    (3, 'ai', ai_command, 2, 'Chat with AI\nUsage: ai <text>'),
    (5, 'aiconfig', ai_config, 5, 'Configure AI\nUsage: aiconfig <key> <value>'),
]

message_act_control = [ai_process_message]