# -*- coding: utf-8 -*-
import os
import logging
import mimetypes
from datetime import datetime
from aiogram import Bot, Dispatcher, types, F
from aiogram.filters import Command
from aiogram.types import FSInputFile
from aiogram.utils.keyboard import InlineKeyboardBuilder
from aiogram.client.session.aiohttp import AiohttpSession
from aiogram.client.telegram import TelegramAPIServer

# === CONFIGURATION ===
# 从环境变量中安全地读取配置
BOT_TOKEN = os.environ.get("BOT_TOKEN")
# 将 USER_ID 转换为整数，如果未设置则默认为 0
AUTHORIZED_USER_ID = int(os.environ.get("AUTHORIZED_USER_ID", 0))
# 允许从环境变量覆盖存储目录，默认为 "local_storage"
# 注意：在 docker-compose 中，这个值被设为 /var/lib/telegram-bot-api
STORAGE_DIR = os.environ.get("STORAGE_DIR", "local_storage")

# === 启动时检查 ===
if not BOT_TOKEN:
    raise ValueError("错误：未设置 BOT_TOKEN 环境变量。")
if AUTHORIZED_USER_ID == 0:
    raise ValueError("错误：未设置 AUTHORIZED_USER_ID 环境变量。")

# === LOGGING ===
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# === INITIAL SETUP ===
local_api_server = TelegramAPIServer.from_base("http://bot-api:8081")
session = AiohttpSession(api=local_api_server)
bot = Bot(token=BOT_TOKEN, session=session)
dp = Dispatcher()

def ensure_storage():
    """Ensure the local storage directory exists."""
    if not os.path.exists(STORAGE_DIR):
        os.makedirs(STORAGE_DIR)
        logger.info(f"Storage folder created at: {STORAGE_DIR}")
    else:
        logger.info(f"Storage folder exists at: {STORAGE_DIR}")

def is_authorized(user_id: int) -> bool:
    """Check if the user is authorized."""
    return user_id == AUTHORIZED_USER_ID

def search_files(keyword: str):
    """Search for files containing the keyword."""
    matches = []
    for root, _, files in os.walk(STORAGE_DIR):
        for f in files:
            if keyword.lower() in f.lower():
                matches.append(os.path.join(root, f))
    return matches

def get_storage_summary():
    """Return summary of all files."""
    total_size = 0
    files_info = []
    for root, dirs, files in os.walk(STORAGE_DIR):
        for f in files:
            path = os.path.join(root, f)
            try:
                size = os.path.getsize(path)
                mtime = datetime.fromtimestamp(os.path.getmtime(path))
                total_size += size
                files_info.append((f, size, mtime))
            except FileNotFoundError:
                logger.warning(f"File not found during summary: {path}")
    return total_size, files_info

def get_extension_from_mime(mime_type: str) -> str:
    """Guess the correct extension from MIME type."""
    ext = mimetypes.guess_extension(mime_type or "")
    return ext if ext else ""

# === COMMAND HANDLERS ===
@dp.message(Command("start"))
async def cmd_start(message: types.Message):
    if not is_authorized(message.from_user.id):
        await message.answer("🚫 Access denied.")
        return
    ensure_storage()
    await message.answer("✅ Bot initialized! Local storage ready.")

@dp.message(Command("overview"))
async def cmd_overview(message: types.Message):
    if not is_authorized(message.from_user.id):
        await message.answer("🚫 Access denied.")
        return

    total_size, files_info = get_storage_summary()
    if not files_info:
        await message.answer("📁 Storage is empty.")
        return

    # 对文件按修改时间降序排序
    files_info.sort(key=lambda x: x[2], reverse=True)
    
    summary = "\n".join(
        [f"{f} — {s/1024:.1f} KB — {d.strftime('%Y-%m-%d %H:%M')}" for f, s, d in files_info]
    )
    await message.answer(f"📦 Files:\n{summary}\n\nTotal: {total_size/1024:.1f} KB")

@dp.message(Command("search"))
async def cmd_search(message: types.Message):
    if not is_authorized(message.from_user.id):
        await message.answer("🚫 Access denied.")
        return

    args = message.text.split(maxsplit=1)
    if len(args) < 2:
        await message.answer("Usage: /search <keyword>")
        return

    keyword = args[1]
    results = search_files(keyword)
    if not results:
        await message.answer("🔍 No matching files found.")
        return

    builder = InlineKeyboardBuilder()
    # 限制搜索结果数量，避免消息过长
    for i, f in enumerate(results[:50]):
        builder.button(text=f"{i+1}. {os.path.basename(f)}", callback_data=f"get|{f}")
    builder.adjust(1) # 每行一个按钮

    await message.answer(f"🔍 Found {len(results)} file(s):", reply_markup=builder.as_markup())

@dp.callback_query(F.data.startswith("get|"))
async def cb_get_file(callback: types.CallbackQuery):
    if not is_authorized(callback.from_user.id):
        await callback.message.answer("🚫 Access denied.")
        return

    path = callback.data.split("|", 1)[1]
    if not os.path.exists(path):
        await callback.message.answer("File not found (it may have been deleted).")
        await callback.answer()
        return
    
    await callback.answer(f"Sending {os.path.basename(path)}...")
    try:
        await callback.message.answer_document(FSInputFile(path))
    except Exception as e:
        logger.error(f"Failed to send file {path}: {e}")
        await callback.message.answer(f"Error sending file: {e}")


@dp.message(Command("delete"))
async def cmd_delete(message: types.Message):
    if not is_authorized(message.from_user.id):
        await message.answer("🚫 Access denied.")
        return

    args = message.text.split(maxsplit=1)
    if len(args) < 2:
        await message.answer("Usage: /delete <filename_keyword>")
        return

    keyword = args[1]
    matches = search_files(keyword)
    if not matches:
        await message.answer("File not found.")
        return

    deleted_files = []
    for f in matches:
        try:
            os.remove(f)
            deleted_files.append(os.path.basename(f))
        except OSError as e:
            logger.error(f"Failed to delete {f}: {e}")

    if deleted_files:
        await message.answer(f"🗑️ Deleted {len(deleted_files)} file(s):\n" + "\n".join(deleted_files))
    else:
        await message.answer("Could not delete matching files.")


# === MESSAGE HANDLER FOR FILES/TEXT ===
@dp.message(F.photo | F.video | F.document | F.text)
async def handle_incoming(message: types.Message):
    if not is_authorized(message.from_user.id):
        await message.answer("🚫 Access denied.")
        return

    ensure_storage()
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

    # --- Handle text ---
    if message.text and not (message.photo or message.video or message.document):
        # ... (文本处理逻辑不变) ...
        filename = f"{timestamp}.txt"
        # 文本文档也放入 'documents' 文件夹
        subdir = "documents"
        target_dir = os.path.join(STORAGE_DIR, subdir)
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
            
        path = os.path.join(target_dir, filename)
        with open(path, "w", encoding="utf-8") as f:
            f.write(message.text)
        await message.answer(f"💾 Text saved as `{filename}` (in /{subdir})", parse_mode="Markdown")
        return

    # --- Handle media/documents ---
    file_obj = None
    file_name = None
    mime_type = None
    is_photo = False

    if message.document:
        file_obj = message.document
        file_name = message.document.file_name
        mime_type = message.document.mime_type
    elif message.photo:
        file_obj = message.photo[-1]  # Highest resolution
        file_name = f"{timestamp}.jpg"
        mime_type = "image/jpeg"
        is_photo = True # 标记为照片
    elif message.video:
        file_obj = message.video
        file_name = message.video.file_name
        mime_type = message.video.mime_type

    if not file_name: # 备用文件名
        file_name = timestamp

    if file_obj:
        logger.info(f"Received file. Original file_name variable: '{file_name}' (MIME: {mime_type}) [IsPhoto: {is_photo}]")
        
        # ... (文件名清理逻辑不变) ...
        base_name, ext = os.path.splitext(file_name)
        if not ext:
            ext = get_extension_from_mime(mime_type)
        if not ext:
            ext = ".bin"  # Fallback
        
        safe_base_name = "".join(c for c in base_name if c.isalnum() or c in (' ', '.', '_', '-')).rstrip()
        final_name = f"{timestamp}_{safe_base_name}{ext}"
        
        # ... (子目录分类逻辑不变) ...
        subdir = "others" # 默认目录
        if mime_type:
            if mime_type.startswith("video/"):
                subdir = "video"
            elif mime_type.startswith("image/"):
                subdir = "pictures"
            elif mime_type in ("application/zip", "application/x-zip-compressed", "application/x-rar-compressed", "application/gzip", "application/x-7z-compressed"):
                subdir = "zip"
            elif mime_type.startswith("audio/"):
                subdir = "audio"
            elif mime_type in ("application/pdf", "application/msword", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "text/plain"):
                subdir = "documents"
        
        target_dir = os.path.join(STORAGE_DIR, subdir)
        
        if not os.path.exists(target_dir):
            try:
                os.makedirs(target_dir)
                logger.info(f"Created subdirectory: {target_dir}")
            except OSError as e:
                logger.error(f"Failed to create subdirectory {target_dir}: {e}")
                target_dir = STORAGE_DIR
        
        destination_path = os.path.join(target_dir, final_name)
        
        logger.info(f"Target file name set to: '{final_name}' in subdir: '{subdir}'")
        
        # --- 根据 is_photo 区分处理 ---
        try:
            if is_photo:
                # --- 照片处理逻辑 ---
                # 照片无法使用 get_file + rename，必须直接下载
                logger.info(f"Processing as Photo. Downloading file_id: {file_obj.file_id}")
                await bot.download_file(file_obj.file_id, destination_path)
                logger.info(f"Downloaded photo to: {destination_path}")

            else:
                # --- 视频/文档处理逻辑 ---
                # 1. 从 API 获取文件信息
                file_info = await bot.get_file(file_obj.file_id)
                
                # 2. 这是文件在 API 服务器上的绝对路径
                source_path = file_info.file_path
                
                logger.info(f"Processing as Document/Video. File is already on disk at: {source_path}")

                # 3. 检查路径
                if not os.path.abspath(source_path).startswith(os.path.abspath(STORAGE_DIR)):
                    logger.error(f"FATAL: File path {source_path} is outside STORAGE_DIR {STORAGE_DIR}. Check docker-compose command.")
                    # 后备方案：尝试下载 (这会很慢，并且可能因404失败)
                    await bot.download_file(file_info.file_path, destination_path)
                else:
                    # 4. 【核心】移动文件
                    os.rename(source_path, destination_path)
                    logger.info(f"Moved file from {source_path} to {destination_path}")
            
            await message.answer(f"💾 Saved `{final_name}` (in /{subdir})", parse_mode="Markdown")
        
        except Exception as e:
            # 确保在这里打印完整的异常，包括 file_id
            logger.error(f"Failed to save file {file_obj.file_id} (Dest: {destination_path}): {e}", exc_info=True)
            await message.answer(f"Error saving file: {e}")

# === MAIN ENTRY ===
async def main():
    ensure_storage()
    logger.info(f"🚀 Bot is starting... Storage: {STORAGE_DIR}")
    await dp.start_polling(bot)

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