# -*- coding: utf-8 -*-
import os
import logging
import mimetypes
import re
import asyncio
from datetime import datetime

from aiogram import Bot, Dispatcher, types, F
from aiogram.filters import Command, StateFilter # <--- [修改1] 引入 StateFilter
from aiogram.types import FSInputFile, ForceReply
from aiogram.utils.keyboard import InlineKeyboardBuilder
from aiogram.client.session.aiohttp import AiohttpSession
from aiogram.client.telegram import TelegramAPIServer
from aiogram.fsm.state import State, StatesGroup
from aiogram.fsm.context import FSMContext

# === CONFIGURATION ===
# 从环境变量中安全地读取配置
BOT_TOKEN = os.environ.get("BOT_TOKEN")
# 将 USER_ID 转换为整数，如果未设置则默认为 0
AUTHORIZED_USER_ID = int(os.environ.get("AUTHORIZED_USER_ID", 0))
# 允许从环境变量覆盖存储目录
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 ===
# 增加 timeout 设置，防止 session 层面的过早断开
local_api_server = TelegramAPIServer.from_base("http://bot-api:8081")
session = AiohttpSession(api=local_api_server, timeout=7200) 
bot = Bot(token=BOT_TOKEN, session=session)
dp = Dispatcher()

# === STATE DEFINITION ===
class RenameState(StatesGroup):
    waiting_for_new_name = State()

# === HELPER FUNCTIONS ===
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:
        pass

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]
    )
    if len(summary) > 3500:
        summary = summary[:3500] + "\n... (list truncated)"
        
    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 ===
# [修改2] 增加 StateFilter(None)，只有在当前没有处于任何 FSM 状态时，才执行这个保存逻辑
# 这样当我们在 waiting_for_new_name 状态时，就不会误触发这里
@dp.message(F.photo | F.video | F.document | F.text, StateFilter(None))
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"
        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

    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"
    elif message.video:
        file_obj = message.video
        file_name = message.video.file_name
        mime_type = message.video.mime_type

    if not file_name: # Fallback name
        file_name = timestamp

    if file_obj:
        logger.info(f"Received file. Original file_name variable: '{file_name}' (MIME: {mime_type})")
        
        base_name, ext = os.path.splitext(file_name)
        if not ext:
            ext = get_extension_from_mime(mime_type)
        if not ext:
            ext = ".bin"
        
        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("image/"):
                subdir = "pictures"
            elif mime_type.startswith("video/"):
                subdir = "video"
            elif mime_type.startswith("audio/"):
                subdir = "audio"
            elif mime_type in ("application/pdf", "application/msword", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/vnd.ms-excel", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "application/vnd.ms-powerpoint", "application/vnd.openxmlformats-officedocument.presentationml.presentation", "application/rtf", "application/epub+zip", "text/plain", "text/csv", "text/markdown", "text/html"):
                subdir = "documents"
            elif mime_type in ("application/zip", "application/vnd.rar", "application/x-zip-compressed", "application/x-rar-compressed", "application/gzip", "application/x-7z-compressed", "application/x-tar", "application/x-bzip2"):
                subdir = "archives"
            elif mime_type in ("application/vnd.android.package-archive", "application/octet-stream", "application/x-msdownload", "application/x-apple-diskimage", "application/x-debian-package", "application/x-rpm", "application/x-sh", "application/x-csh", "application/bat"):
                subdir = "executables"
            elif mime_type in ("application/x-iso9660-image",):
                subdir = "disk_images"
        
        target_dir = os.path.join(STORAGE_DIR, subdir)
        if not os.path.exists(target_dir):
            try:
                os.makedirs(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: '{final_name}' in '{subdir}'")
        
        # === 重试循环 ===
        max_retries = 3
        for attempt in range(max_retries):
            try:
                logger.info(f"Requesting file info for file_id: {file_obj.file_id} (Attempt {attempt+1}/{max_retries})")
                
                # 7200秒超时
                file_info = await bot.get_file(file_obj.file_id, request_timeout=7200)
                
                source_path = file_info.file_path
                logger.info(f"File info received. Path on disk: {source_path}")

                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. Fallback download.")
                    await bot.download_file(file_info.file_path, destination_path)
                else:
                    if os.path.exists(source_path):
                        os.rename(source_path, destination_path)
                        logger.info(f"Moved file to {destination_path}")
                    else:
                        raise FileNotFoundError(f"API says file is at {source_path} but it is missing.")
                
                break
            
            except Exception as e:
                logger.warning(f"Attempt {attempt+1} failed: {e}")
                
                if attempt == max_retries - 1:
                    logger.error("All retry attempts failed.")
                    await message.answer(f"❌ Failed after {max_retries} attempts. Server Error: {e}")
                    return 
                
                await asyncio.sleep(5) 
        
        # === 文件保存成功后 ===
        try:
            builder = InlineKeyboardBuilder()
            builder.button(text="✏️ 重命名", callback_data="rename_file")

            await message.answer(
                f"💾 Saved `{final_name}` (in /{subdir})", 
                parse_mode="Markdown",
                reply_markup=builder.as_markup()
            )
        except Exception as e:
             logger.error(f"Failed to send success message: {e}")


# === RENAME FEATURE HANDLERS ===

@dp.callback_query(F.data == "rename_file")
async def click_rename(callback: types.CallbackQuery, state: FSMContext):
    if not is_authorized(callback.from_user.id):
        return

    message_text = callback.message.text
    match = re.search(r"Saved (.+) \(in /(.+)\)", message_text)
    
    if not match:
        await callback.answer("❌ 无法解析文件路径，可能消息格式已变。", show_alert=True)
        return

    old_filename = match.group(1).strip()
    subdir = match.group(2).strip()
    
    full_path = os.path.join(STORAGE_DIR, subdir, old_filename)
    
    if not os.path.exists(full_path):
        await callback.answer("❌ 文件已不存在（可能已被删除或移动）。", show_alert=True)
        return

    await state.update_data(
        file_path=full_path, 
        file_dir=os.path.dirname(full_path), 
        old_ext=os.path.splitext(old_filename)[1]
    )
    
    await state.set_state(RenameState.waiting_for_new_name)
    
    await callback.message.answer(
        f"当前文件名: `{old_filename}`\n请直接回复**新的文件名** (无需后缀):", 
        parse_mode="Markdown",
        reply_markup=ForceReply(selective=True)
    )
    await callback.answer()

@dp.message(RenameState.waiting_for_new_name)
async def process_rename(message: types.Message, state: FSMContext):
    if not is_authorized(message.from_user.id):
        return

    new_name_input = message.text.strip()
    safe_new_name = "".join(c for c in new_name_input if c.isalnum() or c in (' ', '.', '_', '-')).rstrip()
    
    if not safe_new_name:
        await message.answer("❌ 文件名无效，请重新输入。")
        return 

    data = await state.get_data()
    old_path = data.get('file_path')
    file_dir = data.get('file_dir')
    old_ext = data.get('old_ext')

    if not old_path or not os.path.exists(old_path):
        await message.answer("❌ 原文件已丢失，操作取消。")
        await state.clear()
        return

    if not os.path.splitext(safe_new_name)[1]:
        safe_new_name += old_ext
        
    new_path = os.path.join(file_dir, safe_new_name)

    if os.path.exists(new_path):
        await message.answer("❌ 目标文件名已存在，请换一个名字。")
        return

    try:
        os.rename(old_path, new_path)
        await message.answer(f"✅ 重命名成功！\n`{safe_new_name}`", parse_mode="Markdown")
    except OSError as e:
        logger.error(f"Rename failed: {e}")
        await message.answer(f"❌ 重命名失败: {e}")
    
    await state.clear()

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

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