import re
import asyncio
import logging
import time
from telegram import Update, Message
from telegram.constants import ChatAction, ChatType
from telegram.ext import ContextTypes
from services import auth, ai_client
from utils import context, chunker, persian, formatter

logger = logging.getLogger(__name__)
_MEDIA_GROUP_WAIT = 1.2


# ── Sending helpers ───────────────────────────────────────────

async def _send(message: Message, text: str):
    """Send text with HTML formatting; fall back to plain on parse error."""
    formatted, mode = formatter.safe_send(text)
    for chunk in chunker.chunk_text(formatted):
        try:
            await message.reply_text(chunk, parse_mode=mode)
        except Exception:
            plain = re.sub(r"<[^>]+>", "", chunk)
            await message.reply_text(plain)


async def _send_images(message: Message, urls: list[str]):
    """Upload up to 5 image URLs to Telegram."""
    sent = 0
    for url in urls[:5]:
        try:
            await message.reply_photo(url)
            sent += 1
        except Exception as exc:
            logger.warning("Could not send image %s: %s", url[:60], exc)
    return sent


# ── Group-mention / reply detection ──────────────────────────

def _is_addressed(message: Message, bot_username: str, bot_id: int) -> bool:
    text = message.text or message.caption or ""

    # Reply to bot's own message
    rt = message.reply_to_message
    if rt and rt.from_user:
        if rt.from_user.id == bot_id:
            return True
        if rt.from_user.username and rt.from_user.username.lower() == bot_username.lower():
            return True

    # @mention in entities
    for entity in list(message.entities or []) + list(message.caption_entities or []):
        if entity.type == "mention":
            mention = text[entity.offset: entity.offset + entity.length]
            if mention.lower() == f"@{bot_username.lower()}":
                return True

    # Name mention (Puppy / پاپی)
    return persian.is_puppy_mentioned(text)


def _clean_text(text: str, bot_username: str) -> str:
    text = re.sub(rf"@{re.escape(bot_username)}\s*", "", text, flags=re.IGNORECASE)
    text = persian.strip_puppy_prefix(text)
    return text.strip()


def _get_quoted_context(message: Message, bot_id: int) -> str:
    """If user replied to another user's message, return a quoted prefix for context."""
    rt = message.reply_to_message
    if not rt or not rt.from_user:
        return ""
    if rt.from_user.id == bot_id:
        return ""  # replying to bot is handled naturally via history
    quoted = rt.text or rt.caption or ""
    if not quoted:
        return ""
    name = rt.from_user.first_name or "someone"
    return f'[Replying to {name}: "{quoted[:300]}"]\n'


def _image_analysis_mode(caption: str) -> str:
    text = caption.lower()
    if any(word in text for word in ("ocr", "text", "read", "extract", "متن", "بخوان")):
        return "ocr"
    if any(word in text for word in ("describe", "description", "توصیف", "شرح")):
        return "describe"
    if any(word in text for word in ("objects", "things", "items", "اشیا", "چیزها")):
        return "objects"
    if "?" in caption or any(word in text for word in ("what", "why", "how", "where", "who", "چی", "چرا", "کجا")):
        return "question"
    return "auto"


async def _download_image(message: Message, ctx: ContextTypes.DEFAULT_TYPE) -> tuple[bytes, str] | None:
    if message.photo:
        image_file = await ctx.bot.get_file(message.photo[-1].file_id)
        return bytes(await image_file.download_as_bytearray()), "image/jpeg"
    if message.document and (message.document.mime_type or "").startswith("image/"):
        image_file = await ctx.bot.get_file(message.document.file_id)
        return bytes(await image_file.download_as_bytearray()), message.document.mime_type or "image/jpeg"
    return None


def _is_image_message(message: Message) -> bool:
    return bool(message.photo or (message.document and (message.document.mime_type or "").startswith("image/")))


# ── Image handler ─────────────────────────────────────────────

async def _handle_photo(update: Update, ctx: ContextTypes.DEFAULT_TYPE, force_mode: str | None = None):
    message = update.effective_message
    user = update.effective_user
    chat = update.effective_chat

    caption = _clean_text(message.caption or "", ctx.bot.username or "")
    settings = context.get_user_settings(user.id)
    history = context.get_history(message.chat_id)
    sys_prompt = context.get_system_prompt()
    if context.reset_chat_if_input_too_large(message.chat_id, sys_prompt, keep_last=False):
        history = []

    await ctx.bot.send_chat_action(chat.id, ChatAction.TYPING)

    try:
        downloaded = await _download_image(message, ctx)
        if not downloaded:
            return

        image_bytes, mime_type = downloaded
        analysis_mode = force_mode or _image_analysis_mode(caption)

        reply = await ai_client.chat_with_image(
            messages=history,
            system_prompt=sys_prompt,
            max_tokens=settings["max_tokens"],
            image_bytes=image_bytes,
            mime_type=mime_type,
            caption=caption,
            analysis_mode=analysis_mode,
        )

        context.add_message(message.chat_id, "user", f"[Photo]{': ' + caption if caption else ''}")
        context.add_message(message.chat_id, "assistant", reply)
        context.increment_msg()
        context.record_event("images_analyzed")

        await _send(message, reply)

    except Exception as exc:
        context.record_error(f"Image error: {exc}")
        await message.reply_text(f"❌ Image error: {exc}")


async def _handle_media_group(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    message = update.effective_message
    group_id = message.media_group_id
    if not group_id:
        await _handle_photo(update, ctx)
        return

    groups = ctx.application.bot_data.setdefault("media_groups", {})
    bucket = groups.setdefault(group_id, {"messages": [], "last_seen": 0.0})
    bucket["messages"].append(message)
    bucket["last_seen"] = time.time()

    if len(bucket["messages"]) > 1:
        return

    await asyncio.sleep(_MEDIA_GROUP_WAIT)
    bucket = groups.pop(group_id, bucket)
    messages = bucket["messages"][:10]
    first = messages[0]
    user = first.from_user
    settings = context.get_user_settings(user.id)
    history = context.get_history(first.chat_id)
    sys_prompt = context.get_system_prompt()
    if context.reset_chat_if_input_too_large(first.chat_id, sys_prompt, keep_last=False):
        history = []
    caption = ""
    images: list[tuple[bytes, str]] = []

    await ctx.bot.send_chat_action(first.chat_id, ChatAction.TYPING)

    try:
        for msg in messages:
            if not caption and msg.caption:
                caption = _clean_text(msg.caption, ctx.bot.username or "")
            downloaded = await _download_image(msg, ctx)
            if downloaded:
                images.append(downloaded)

        if not images:
            return

        reply = await ai_client.chat_with_images(
            messages=history,
            system_prompt=sys_prompt,
            max_tokens=settings["max_tokens"],
            images=images,
            caption=caption,
            analysis_mode=_image_analysis_mode(caption),
        )

        context.add_message(first.chat_id, "user", f"[{len(images)} photos]{': ' + caption if caption else ''}")
        context.add_message(first.chat_id, "assistant", reply)
        context.increment_msg()
        context.record_event("multi_image_analyses")

        await _send(first, reply)

    except Exception as exc:
        context.record_error(f"Multi-image error: {exc}")
        await first.reply_text(f"❌ Multi-image error: {exc}")


# ── Main message handler ──────────────────────────────────────

async def message_handler(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    message = update.effective_message
    if not message:
        return

    user = update.effective_user
    chat = update.effective_chat

    if not auth.is_authorized(user.id, user.username or ""):
        return
    auth.register_known(user.id, username=user.username or "")

    # Group gate: must be addressed
    if chat.type in (ChatType.GROUP, ChatType.SUPERGROUP):
        if not _is_addressed(message, ctx.bot.username or "", ctx.bot.id):
            return

    # Image branch: Telegram photos and image files/documents.
    if _is_image_message(message):
        if message.media_group_id:
            await _handle_media_group(update, ctx)
        else:
            await _handle_photo(update, ctx)
        return

    text = message.text or ""
    if not text:
        return

    text = _clean_text(text, ctx.bot.username or "")
    if not text:
        return

    # Include quoted message from another user as context prefix
    quoted = _get_quoted_context(message, ctx.bot.id)
    full_text = quoted + text

    settings = context.get_user_settings(user.id)
    sys_prompt = context.get_system_prompt()

    # Add to history first so AI sees the current message
    context.add_message(message.chat_id, "user", full_text)
    context.reset_chat_if_input_too_large(message.chat_id, sys_prompt, keep_last=True)
    history = context.get_history(message.chat_id)

    await ctx.bot.send_chat_action(chat.id, ChatAction.TYPING)

    try:
        reply, image_urls = await ai_client.chat(
            messages=history,
            system_prompt=sys_prompt,
            max_tokens=settings["max_tokens"],
            search_enabled=settings["search_enabled"],
        )

        context.add_message(message.chat_id, "assistant", reply)
        context.increment_msg()
        context.record_event("text_chats")

        # Send images first (appear above text)
        if image_urls:
            await _send_images(message, image_urls)

        await _send(message, reply)

    except Exception as exc:
        context.pop_last_message(message.chat_id)
        context.record_error(f"Chat error: {exc}")
        logger.error("Chat error: %s", exc)
        await message.reply_text(f"❌ Error: {exc}")