#!/usr/bin/env python3
"""
Telegram Data Collector для OpenClaw
-----------------------------------
Скрипт для сбора сообщений из Telegram-групп и каналов
и сохранения их в формате, удобном для анализа с помощью OpenClaw.
"""

import os
import sys
import json
import asyncio
from datetime import datetime, timedelta
from pathlib import Path
from telethon import TelegramClient, events
from telethon.tl.functions.messages import GetHistoryRequest
import logging
from dotenv import load_dotenv

# Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.path.join(SKILL_DIR, "logs", "telegram_collector.log")),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger(__name__)

# Путь к директории скилла
SKILL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# Загрузка переменных окружения из .env файла
load_dotenv(os.path.expanduser("~/.openclaw/.env"))

# Конфигурация
API_ID = os.getenv('TELEGRAM_COLLECTOR_API_ID')
API_HASH = os.getenv('TELEGRAM_COLLECTOR_API_HASH')
PHONE_NUMBER = os.getenv('TELEGRAM_COLLECTOR_PHONE')
SESSION_FILE = os.getenv('TELEGRAM_COLLECTOR_SESSION', 'telegram_collector')
OUTPUT_DIR = os.getenv('OUTPUT_DIR', os.path.join(SKILL_DIR, '..', 'data', 'telegram-collector'))

# Создание директорий для выходных данных
Path(OUTPUT_DIR).mkdir(exist_ok=True)
Path(f"{OUTPUT_DIR}/raw").mkdir(exist_ok=True)
Path(f"{OUTPUT_DIR}/topics").mkdir(exist_ok=True)
Path(f"{OUTPUT_DIR}/summaries").mkdir(exist_ok=True)

class TelegramCollector:
    def __init__(self):
        self.client = None
        self.groups = {}
        self.config_file = os.path.join(SKILL_DIR, 'groups_config.json')
        self.load_config()

    def load_config(self):
        """Загружает конфигурацию групп из JSON-файла."""
        try:
            if os.path.exists(self.config_file):
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    self.groups = json.load(f)
                logger.info(f"Конфигурация загружена: {len(self.groups)} групп/каналов")
            else:
                logger.warning(f"Файл конфигурации {self.config_file} не найден. Создаем новый.")
                self.groups = {}
                self.save_config()
        except Exception as e:
            logger.error(f"Ошибка загрузки конфигурации: {e}")
            self.groups = {}

    def save_config(self):
        """Сохраняет текущую конфигурацию групп в JSON-файл."""
        with open(self.config_file, 'w', encoding='utf-8') as f:
            json.dump(self.groups, f, indent=2, ensure_ascii=False)
        logger.info("Конфигурация сохранена")

    def update_group_config(self, group_id, name, username=None, topics=None, last_update=None):
        """Обновляет конфигурацию для указанной группы."""
        if str(group_id) not in self.groups:
            self.groups[str(group_id)] = {
                "name": name,
                "username": username,
                "topics": topics or [],
                "last_update": last_update or "1970-01-01 00:00:00"
            }
        else:
            if name:
                self.groups[str(group_id)]["name"] = name
            if username:
                self.groups[str(group_id)]["username"] = username
            if topics:
                self.groups[str(group_id)]["topics"] = topics
            if last_update:
                self.groups[str(group_id)]["last_update"] = last_update
        
        self.save_config()

    async def connect(self):
        """Подключается к Telegram API."""
        if not all([API_ID, API_HASH, PHONE_NUMBER]):
            logger.error("API_ID, API_HASH или PHONE_NUMBER не установлены. Проверьте .env файл.")
            sys.exit(1)
            
        try:
            self.client = TelegramClient(SESSION_FILE, API_ID, API_HASH)
            await self.client.start(phone=PHONE_NUMBER)
            logger.info("Успешное подключение к Telegram API")
            return True
        except Exception as e:
            logger.error(f"Ошибка подключения к Telegram API: {e}")
            return False

    async def get_dialogs(self):
        """Получает список доступных диалогов."""
        dialogs = await self.client.get_dialogs()
        groups_channels = [d for d in dialogs if d.is_group or d.is_channel]
        
        logger.info(f"Доступно {len(groups_channels)} групп и каналов:")
        for i, dialog in enumerate(groups_channels, 1):
            dialog_type = "канал" if dialog.is_channel else "группа"
            entity = dialog.entity
            username = getattr(entity, 'username', None)
            logger.info(f"{i}. {dialog.name} (ID: {dialog.id}, @{username}, тип: {dialog_type})")
        
        return groups_channels

    async def get_messages(self, chat_entity, limit=100, since_date=None):
        """Получает сообщения из указанного чата."""
        try:
            # Если указана дата, получаем сообщения только с этой даты
            if since_date:
                since_date = datetime.strptime(since_date, "%Y-%m-%d %H:%M:%S")
                messages = []
                offset_id = 0
                
                while True:
                    history = await self.client(GetHistoryRequest(
                        peer=chat_entity,
                        offset_id=offset_id,
                        offset_date=None,
                        add_offset=0,
                        limit=100,
                        max_id=0,
                        min_id=0,
                        hash=0
                    ))
                    
                    if not history.messages:
                        break
                        
                    for message in history.messages:
                        if message.date < since_date:
                            # Достигли сообщений старше указанной даты, останавливаемся
                            return messages
                        
                        if message.message:  # Пропускаем пустые сообщения
                            messages.append(message)
                    
                    offset_id = history.messages[-1].id
                    if len(messages) >= limit:
                        break
                
                return messages
            else:
                # Если дата не указана, просто берем последние сообщения
                return await self.client.get_messages(chat_entity, limit=limit)
        except Exception as e:
            logger.error(f"Ошибка при получении сообщений: {e}")
            return []

    async def collect_from_group(self, group_id, limit=100):
        """Собирает сообщения из указанной группы."""
        try:
            group_id = int(group_id)
            group_config = self.groups.get(str(group_id), {})
            last_update = group_config.get("last_update", "1970-01-01 00:00:00")
            
            entity = await self.client.get_entity(group_id)
            logger.info(f"Сбор данных из {entity.title} (ID: {group_id})")
            
            messages = await self.get_messages(entity, limit=limit, since_date=last_update)
            logger.info(f"Получено {len(messages)} новых сообщений")
            
            # Обработка и сохранение сообщений
            if messages:
                processed_messages = []
                for msg in messages:
                    processed_msg = {
                        "id": msg.id,
                        "date": msg.date.strftime("%Y-%m-%d %H:%M:%S"),
                        "from_id": getattr(msg.from_id, "user_id", None) if msg.from_id else None,
                        "text": msg.message,
                        "has_media": bool(msg.media),
                        "reply_to_msg_id": msg.reply_to_msg_id
                    }
                    processed_messages.append(processed_msg)
                
                # Сохраняем сырые данные
                today = datetime.now().strftime("%Y-%m-%d")
                raw_file = f"{OUTPUT_DIR}/raw/{entity.id}_{today}.json"
                
                with open(raw_file, 'w', encoding='utf-8') as f:
                    json.dump({
                        "group_id": group_id,
                        "group_name": entity.title,
                        "collected_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                        "messages_count": len(processed_messages),
                        "messages": processed_messages
                    }, f, indent=2, ensure_ascii=False)
                
                # Также сохраняем в формате Markdown для удобного просмотра
                md_file = f"{OUTPUT_DIR}/raw/{entity.id}_{today}.md"
                with open(md_file, 'w', encoding='utf-8') as f:
                    f.write(f"# Сырые данные из {entity.title}\n\n")
                    f.write(f"## Собрано: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                    f.write(f"## Всего сообщений: {len(processed_messages)}\n\n")
                    
                    for i, msg in enumerate(processed_messages, 1):
                        f.write(f"### Сообщение {i}\n")
                        f.write(f"**ID**: {msg['id']}\n")
                        f.write(f"**Дата**: {msg['date']}\n")
                        f.write(f"**Текст**:\n{msg['text']}\n")
                        f.write(f"**Содержит медиа**: {'Да' if msg['has_media'] else 'Нет'}\n\n")
                
                # Обновляем конфигурацию группы с новой датой последнего обновления
                if processed_messages:
                    latest_date = processed_messages[0]["date"]
                    self.update_group_config(
                        group_id, 
                        name=entity.title,
                        username=getattr(entity, 'username', None),
                        last_update=latest_date
                    )
                
                return len(processed_messages)
            else:
                logger.info(f"В группе {entity.title} нет новых сообщений с {last_update}")
                return 0
        except Exception as e:
            logger.error(f"Ошибка при сборе данных из группы {group_id}: {e}")
            return 0

    async def list_groups(self):
        """Выводит список групп в конфигурации."""
        if not self.groups:
            logger.info("В конфигурации нет сохраненных групп")
            return
        
        logger.info("Сохраненные группы и каналы:")
        for group_id, info in self.groups.items():
            logger.info(f"ID: {group_id}, Название: {info['name']}, Последнее обновление: {info['last_update']}")

    async def add_group(self, group_id, topics=None):
        """Добавляет группу в список отслеживаемых."""
        try:
            group_id = int(group_id)
            entity = await self.client.get_entity(group_id)
            
            self.update_group_config(
                group_id,
                name=entity.title,
                username=getattr(entity, 'username', None),
                topics=topics
            )
            
            logger.info(f"Группа {entity.title} (ID: {group_id}) добавлена в конфигурацию")
            return True
        except Exception as e:
            logger.error(f"Ошибка при добавлении группы {group_id}: {e}")
            return False

    async def remove_group(self, group_id):
        """Удаляет группу из списка отслеживаемых."""
        if str(group_id) in self.groups:
            group_name = self.groups[str(group_id)]["name"]
            del self.groups[str(group_id)]
            self.save_config()
            logger.info(f"Группа {group_name} (ID: {group_id}) удалена из конфигурации")
            return True
        else:
            logger.warning(f"Группа с ID {group_id} не найдена в конфигурации")
            return False

    async def collect_all(self, limit=100):
        """Собирает данные из всех групп в конфигурации."""
        total_messages = 0
        
        for group_id in self.groups:
            count = await self.collect_from_group(group_id, limit=limit)
            total_messages += count
        
        logger.info(f"Всего собрано {total_messages} новых сообщений из {len(self.groups)} групп/каналов")
        return total_messages

async def main():
    """Основная функция для запуска коллектора."""
    collector = TelegramCollector()
    
    if len(sys.argv) < 2:
        print("Использование: python collector.py <команда> [аргументы]")
        print("Команды:")
        print("  list - вывести список доступных диалогов")
        print("  list-config - вывести список сохраненных групп")
        print("  add <group_id> [topic1,topic2,...] - добавить группу")
        print("  remove <group_id> - удалить группу")
        print("  collect <group_id> [limit] - собрать данные из группы")
        print("  collect-all [limit] - собрать данные из всех групп")
        sys.exit(1)
    
    command = sys.argv[1].lower()
    
    if not await collector.connect():
        sys.exit(1)
    
    try:
        if command == "list":
            await collector.get_dialogs()
        
        elif command == "list-config":
            await collector.list_groups()
        
        elif command == "add" and len(sys.argv) >= 3:
            group_id = sys.argv[2]
            topics = sys.argv[3].split(",") if len(sys.argv) > 3 else None
            await collector.add_group(group_id, topics)
        
        elif command == "remove" and len(sys.argv) >= 3:
            group_id = sys.argv[2]
            await collector.remove_group(group_id)
        
        elif command == "collect" and len(sys.argv) >= 3:
            group_id = sys.argv[2]
            limit = int(sys.argv[3]) if len(sys.argv) > 3 else 100
            await collector.collect_from_group(group_id, limit)
        
        elif command == "collect-all":
            limit = int(sys.argv[2]) if len(sys.argv) > 2 else 100
            await collector.collect_all(limit)
        
        else:
            print("Неизвестная команда или недостаточно аргументов")
            sys.exit(1)
    
    finally:
        await collector.client.disconnect()

if __name__ == "__main__":
    asyncio.run(main())