"""
Telegram MCP Server
Enables AI agents to interact with Telegram via a Bot Token.
Supports sending/receiving text, multimedia, and file downloads.
"""

import os
import json
import base64
import httpx
from typing import Optional, List
from pydantic import BaseModel, Field, ConfigDict
from mcp.server.fastmcp import FastMCP

# Load .env if present (no-op if not found)
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

# ─── Server Init ─────────────────────────────────────────────────────────────

mcp = FastMCP("telegram_mcp")

BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN", "")
BASE_URL = f"https://api.telegram.org/bot{BOT_TOKEN}"
FILE_URL = f"https://api.telegram.org/file/bot{BOT_TOKEN}"

# ─── Shared Helpers ───────────────────────────────────────────────────────────

async def _tg(method: str, payload: dict) -> dict:
    """Make a Telegram Bot API request."""
    if not BOT_TOKEN:
        raise ValueError("TELEGRAM_BOT_TOKEN environment variable is not set.")
    async with httpx.AsyncClient(timeout=60) as client:
        r = await client.post(f"{BASE_URL}/{method}", json=payload)
        r.raise_for_status()
        return r.json()

async def _tg_multipart(method: str, data: dict, files: dict) -> dict:
    """Make a multipart/form-data Telegram request (for file uploads)."""
    if not BOT_TOKEN:
        raise ValueError("TELEGRAM_BOT_TOKEN environment variable is not set.")
    async with httpx.AsyncClient(timeout=120) as client:
        r = await client.post(f"{BASE_URL}/{method}", data=data, files=files)
        r.raise_for_status()
        return r.json()

def _ok(result: dict) -> str:
    """Format a successful Telegram response as JSON."""
    if not result.get("ok"):
        return json.dumps({"error": result.get("description", "Unknown Telegram error")}, indent=2)
    return json.dumps(result.get("result", {}), indent=2)

def _err(e: Exception) -> str:
    """Format an error response."""
    if isinstance(e, httpx.HTTPStatusError):
        try:
            body = e.response.json()
            return json.dumps({"error": body.get("description", str(e))}, indent=2)
        except Exception:
            pass
        return json.dumps({"error": f"HTTP {e.response.status_code}: {e.response.text}"}, indent=2)
    if isinstance(e, httpx.TimeoutException):
        return json.dumps({"error": "Request timed out. Try again."}, indent=2)
    return json.dumps({"error": str(e)}, indent=2)


# ─── Input Models ─────────────────────────────────────────────────────────────

class SendMessageInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username (e.g. '123456789' or '@mychannel')")
    text: str = Field(..., description="Message text (supports Markdown/HTML via parse_mode)", min_length=1, max_length=4096)
    parse_mode: Optional[str] = Field(default=None, description="Formatting: 'Markdown', 'MarkdownV2', or 'HTML'")
    reply_to_message_id: Optional[int] = Field(default=None, description="Reply to a specific message ID")
    disable_notification: Optional[bool] = Field(default=False, description="Send silently (no notification sound)")

class SendPhotoInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    photo: str = Field(..., description="HTTP URL of the photo, a Telegram file_id, or a local file path")
    caption: Optional[str] = Field(default=None, description="Photo caption (max 1024 chars)", max_length=1024)
    parse_mode: Optional[str] = Field(default=None, description="Caption formatting: 'Markdown' or 'HTML'")
    reply_to_message_id: Optional[int] = Field(default=None, description="Reply to a specific message ID")

class SendVideoInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    video: str = Field(..., description="HTTP URL of the video, a Telegram file_id, or a local file path")
    caption: Optional[str] = Field(default=None, description="Video caption (max 1024 chars)", max_length=1024)
    parse_mode: Optional[str] = Field(default=None, description="Caption formatting: 'Markdown' or 'HTML'")
    width: Optional[int] = Field(default=None, description="Video width in pixels")
    height: Optional[int] = Field(default=None, description="Video height in pixels")
    duration: Optional[int] = Field(default=None, description="Video duration in seconds")

class SendAudioInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    audio: str = Field(..., description="HTTP URL of the audio, a Telegram file_id, or a local file path")
    caption: Optional[str] = Field(default=None, description="Audio caption", max_length=1024)
    duration: Optional[int] = Field(default=None, description="Audio duration in seconds")
    performer: Optional[str] = Field(default=None, description="Audio performer name")
    title: Optional[str] = Field(default=None, description="Audio track title")

class SendDocumentInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    document: str = Field(..., description="HTTP URL of the document, a Telegram file_id, or a local file path")
    caption: Optional[str] = Field(default=None, description="Document caption", max_length=1024)
    parse_mode: Optional[str] = Field(default=None, description="Caption formatting: 'Markdown' or 'HTML'")

class SendVoiceInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    voice: str = Field(..., description="HTTP URL of the voice file (.ogg), a Telegram file_id, or a local file path")
    caption: Optional[str] = Field(default=None, description="Voice message caption", max_length=1024)
    duration: Optional[int] = Field(default=None, description="Voice message duration in seconds")

class SendStickerInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    sticker: str = Field(..., description="Telegram file_id of the sticker or HTTP URL (.webp)")

class SendLocationInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    chat_id: str = Field(..., description="Target chat ID or @username")
    latitude: float = Field(..., description="Latitude of the location", ge=-90, le=90)
    longitude: float = Field(..., description="Longitude of the location", ge=-180, le=180)

class GetUpdatesInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    offset: Optional[int] = Field(default=None, description="Offset update ID to fetch updates after (marks previous as read)")
    limit: Optional[int] = Field(default=20, description="Max number of updates to retrieve (1-100)", ge=1, le=100)
    allowed_updates: Optional[List[str]] = Field(default=None, description="Types to filter: ['message','callback_query', etc.]")

class GetChatInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Chat ID or @username to get info about")

class GetFileInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    file_id: str = Field(..., description="Telegram file_id to get download info for")

class DownloadFileInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    file_id: str = Field(..., description="Telegram file_id to download")
    save_path: Optional[str] = Field(default=None, description="Local path to save the file. If omitted, returns base64-encoded content.")

class ForwardMessageInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Target chat to forward to")
    from_chat_id: str = Field(..., description="Source chat ID where the original message is")
    message_id: int = Field(..., description="Message ID to forward")

class DeleteMessageInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    chat_id: str = Field(..., description="Chat ID containing the message")
    message_id: int = Field(..., description="Message ID to delete")

class PinMessageInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    chat_id: str = Field(..., description="Chat ID")
    message_id: int = Field(..., description="Message ID to pin")
    disable_notification: Optional[bool] = Field(default=False, description="Pin silently without notification")

class GetChatMembersCountInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    chat_id: str = Field(..., description="Chat ID or @username")


# ─── Helpers for file sending ─────────────────────────────────────────────────

def _is_local_path(value: str) -> bool:
    return os.path.exists(value)

async def _send_media(method: str, chat_id: str, media_field: str,
                      media_value: str, extra: dict) -> str:
    """Dispatch to URL/file_id send or multipart upload."""
    try:
        if _is_local_path(media_value):
            filename = os.path.basename(media_value)
            with open(media_value, "rb") as f:
                content = f.read()
            data = {"chat_id": chat_id, **{k: v for k, v in extra.items() if v is not None}}
            files = {media_field: (filename, content)}
            result = await _tg_multipart(method, data, files)
        else:
            payload = {"chat_id": chat_id, media_field: media_value,
                       **{k: v for k, v in extra.items() if v is not None}}
            result = await _tg(method, payload)
        return _ok(result)
    except Exception as e:
        return _err(e)


# ─── Tools ────────────────────────────────────────────────────────────────────

@mcp.tool(
    name="telegram_send_message",
    annotations={
        "title": "Send Text Message",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_message(params: SendMessageInput) -> str:
    """Send a text message to a Telegram chat.

    Supports plain text, Markdown, and HTML formatting.
    Can reply to an existing message using reply_to_message_id.

    Args:
        params (SendMessageInput): Validated input with chat_id, text, parse_mode,
            reply_to_message_id, and disable_notification.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    try:
        payload: dict = {
            "chat_id": params.chat_id,
            "text": params.text,
        }
        if params.parse_mode:
            payload["parse_mode"] = params.parse_mode
        if params.reply_to_message_id:
            payload["reply_to_message_id"] = params.reply_to_message_id
        if params.disable_notification:
            payload["disable_notification"] = params.disable_notification
        result = await _tg("sendMessage", payload)
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_send_photo",
    annotations={
        "title": "Send Photo",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_photo(params: SendPhotoInput) -> str:
    """Send a photo to a Telegram chat.

    Accepts a local file path, a public HTTP(S) URL, or a Telegram file_id.
    Optionally include a caption (supports Markdown/HTML).

    Args:
        params (SendPhotoInput): chat_id, photo source, optional caption/parse_mode.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    extra = {
        "caption": params.caption,
        "parse_mode": params.parse_mode,
        "reply_to_message_id": params.reply_to_message_id,
    }
    return await _send_media("sendPhoto", params.chat_id, "photo", params.photo, extra)


@mcp.tool(
    name="telegram_send_video",
    annotations={
        "title": "Send Video",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_video(params: SendVideoInput) -> str:
    """Send a video to a Telegram chat.

    Accepts a local file path, a public HTTP(S) URL, or a Telegram file_id.

    Args:
        params (SendVideoInput): chat_id, video source, optional caption, dimensions.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    extra = {
        "caption": params.caption,
        "parse_mode": params.parse_mode,
        "width": params.width,
        "height": params.height,
        "duration": params.duration,
    }
    return await _send_media("sendVideo", params.chat_id, "video", params.video, extra)


@mcp.tool(
    name="telegram_send_audio",
    annotations={
        "title": "Send Audio",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_audio(params: SendAudioInput) -> str:
    """Send an audio file to a Telegram chat (shown as music player).

    Accepts a local file path, a public HTTP(S) URL, or a Telegram file_id.

    Args:
        params (SendAudioInput): chat_id, audio source, optional caption/performer/title.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    extra = {
        "caption": params.caption,
        "duration": params.duration,
        "performer": params.performer,
        "title": params.title,
    }
    return await _send_media("sendAudio", params.chat_id, "audio", params.audio, extra)


@mcp.tool(
    name="telegram_send_document",
    annotations={
        "title": "Send Document / File",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_document(params: SendDocumentInput) -> str:
    """Send any file (document) to a Telegram chat.

    Use this for PDFs, ZIPs, code files, or any file type that doesn't fit
    photo/video/audio. Accepts a local path, URL, or Telegram file_id.

    Args:
        params (SendDocumentInput): chat_id, document source, optional caption.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    extra = {
        "caption": params.caption,
        "parse_mode": params.parse_mode,
    }
    return await _send_media("sendDocument", params.chat_id, "document", params.document, extra)


@mcp.tool(
    name="telegram_send_voice",
    annotations={
        "title": "Send Voice Message",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_voice(params: SendVoiceInput) -> str:
    """Send a voice message (.ogg/Opus) to a Telegram chat.

    Accepts a local file path, a public HTTP(S) URL, or a Telegram file_id.

    Args:
        params (SendVoiceInput): chat_id, voice file source, optional caption/duration.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    extra = {
        "caption": params.caption,
        "duration": params.duration,
    }
    return await _send_media("sendVoice", params.chat_id, "voice", params.voice, extra)


@mcp.tool(
    name="telegram_send_sticker",
    annotations={
        "title": "Send Sticker",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_sticker(params: SendStickerInput) -> str:
    """Send a sticker to a Telegram chat using a file_id or .webp URL.

    Args:
        params (SendStickerInput): chat_id and sticker file_id or URL.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    try:
        result = await _tg("sendSticker", {"chat_id": params.chat_id, "sticker": params.sticker})
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_send_location",
    annotations={
        "title": "Send Location",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_send_location(params: SendLocationInput) -> str:
    """Send a geographic location (map pin) to a Telegram chat.

    Args:
        params (SendLocationInput): chat_id, latitude, longitude.

    Returns:
        str: JSON with the sent Message object or an error description.
    """
    try:
        result = await _tg("sendLocation", {
            "chat_id": params.chat_id,
            "latitude": params.latitude,
            "longitude": params.longitude,
        })
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_get_updates",
    annotations={
        "title": "Get Updates (Messages)",
        "readOnlyHint": True,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": True
    }
)
async def telegram_get_updates(params: GetUpdatesInput) -> str:
    """Retrieve recent updates (messages, callbacks, etc.) sent to the bot.

    Use the 'offset' parameter to mark previous updates as read and paginate.
    Call with offset = (last_update_id + 1) to acknowledge processed messages.

    Args:
        params (GetUpdatesInput): offset, limit, allowed_updates filter list.

    Returns:
        str: JSON list of Update objects containing messages and metadata.
    """
    try:
        payload: dict = {"limit": params.limit}
        if params.offset is not None:
            payload["offset"] = params.offset
        if params.allowed_updates:
            payload["allowed_updates"] = params.allowed_updates
        result = await _tg("getUpdates", payload)
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_get_chat_info",
    annotations={
        "title": "Get Chat Info",
        "readOnlyHint": True,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": True
    }
)
async def telegram_get_chat_info(params: GetChatInput) -> str:
    """Get information about a chat (group, channel, or user).

    Returns name, type, description, member count, and more.

    Args:
        params (GetChatInput): chat_id or @username.

    Returns:
        str: JSON with Chat object details.
    """
    try:
        result = await _tg("getChat", {"chat_id": params.chat_id})
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_get_chat_member_count",
    annotations={
        "title": "Get Chat Member Count",
        "readOnlyHint": True,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": True
    }
)
async def telegram_get_chat_member_count(params: GetChatMembersCountInput) -> str:
    """Get the total number of members in a chat/channel/group.

    Args:
        params (GetChatMembersCountInput): chat_id or @username.

    Returns:
        str: JSON with {"member_count": N}.
    """
    try:
        result = await _tg("getChatMemberCount", {"chat_id": params.chat_id})
        return json.dumps({"member_count": result.get("result")}, indent=2)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_get_file_info",
    annotations={
        "title": "Get File Info",
        "readOnlyHint": True,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": True
    }
)
async def telegram_get_file_info(params: GetFileInput) -> str:
    """Get metadata and download path for a file stored on Telegram servers.

    Returns the file_path needed to construct the download URL:
    https://api.telegram.org/file/bot<TOKEN>/<file_path>

    Args:
        params (GetFileInput): file_id of the Telegram file.

    Returns:
        str: JSON with File object including file_path, file_size, and download_url.
    """
    try:
        result = await _tg("getFile", {"file_id": params.file_id})
        if result.get("ok"):
            file_obj = result["result"]
            file_obj["download_url"] = f"{FILE_URL}/{file_obj.get('file_path', '')}"
            return json.dumps(file_obj, indent=2)
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_download_file",
    annotations={
        "title": "Download File from Telegram",
        "readOnlyHint": True,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": True
    }
)
async def telegram_download_file(params: DownloadFileInput) -> str:
    """Download a file from Telegram using its file_id.

    If save_path is provided, writes the file to disk and returns the path.
    Otherwise returns the file content as base64-encoded string.

    Use telegram_get_file_info first if you only need the download URL.

    Args:
        params (DownloadFileInput): file_id, optional save_path.

    Returns:
        str: JSON with saved_path or base64 content + filename and size.
    """
    try:
        # Step 1: resolve file path
        meta = await _tg("getFile", {"file_id": params.file_id})
        if not meta.get("ok"):
            return _ok(meta)
        file_obj = meta["result"]
        file_path = file_obj.get("file_path", "")
        download_url = f"{FILE_URL}/{file_path}"
        filename = os.path.basename(file_path)

        # Step 2: download content
        async with httpx.AsyncClient(timeout=120) as client:
            r = await client.get(download_url)
            r.raise_for_status()
            content = r.content

        # Step 3: save or return
        if params.save_path:
            dest = params.save_path
            if os.path.isdir(dest):
                dest = os.path.join(dest, filename)
            with open(dest, "wb") as f:
                f.write(content)
            return json.dumps({
                "saved_path": dest,
                "filename": filename,
                "file_size": len(content),
            }, indent=2)
        else:
            return json.dumps({
                "filename": filename,
                "file_size": len(content),
                "content_base64": base64.b64encode(content).decode(),
            }, indent=2)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_forward_message",
    annotations={
        "title": "Forward Message",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_forward_message(params: ForwardMessageInput) -> str:
    """Forward a message from one chat to another.

    Args:
        params (ForwardMessageInput): target chat_id, from_chat_id, message_id.

    Returns:
        str: JSON with the forwarded Message object or an error description.
    """
    try:
        result = await _tg("forwardMessage", {
            "chat_id": params.chat_id,
            "from_chat_id": params.from_chat_id,
            "message_id": params.message_id,
        })
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_delete_message",
    annotations={
        "title": "Delete Message",
        "readOnlyHint": False,
        "destructiveHint": True,
        "idempotentHint": False,
        "openWorldHint": True
    }
)
async def telegram_delete_message(params: DeleteMessageInput) -> str:
    """Delete a message from a chat (bot must have delete permission).

    Args:
        params (DeleteMessageInput): chat_id and message_id to delete.

    Returns:
        str: JSON {"deleted": true} on success or an error description.
    """
    try:
        result = await _tg("deleteMessage", {
            "chat_id": params.chat_id,
            "message_id": params.message_id,
        })
        if result.get("ok"):
            return json.dumps({"deleted": True}, indent=2)
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_pin_message",
    annotations={
        "title": "Pin Message",
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": True
    }
)
async def telegram_pin_message(params: PinMessageInput) -> str:
    """Pin a message in a chat (bot must be admin with pin_messages permission).

    Args:
        params (PinMessageInput): chat_id, message_id, disable_notification flag.

    Returns:
        str: JSON {"pinned": true} on success or an error description.
    """
    try:
        result = await _tg("pinChatMessage", {
            "chat_id": params.chat_id,
            "message_id": params.message_id,
            "disable_notification": params.disable_notification,
        })
        if result.get("ok"):
            return json.dumps({"pinned": True}, indent=2)
        return _ok(result)
    except Exception as e:
        return _err(e)


@mcp.tool(
    name="telegram_get_bot_info",
    annotations={
        "title": "Get Bot Info",
        "readOnlyHint": True,
        "destructiveHint": False,
        "idempotentHint": True,
        "openWorldHint": False
    }
)
async def telegram_get_bot_info() -> str:
    """Get information about the configured bot (name, username, capabilities).

    Useful to verify the bot token is valid and check bot identity.

    Returns:
        str: JSON with User object representing the bot.
    """
    try:
        result = await _tg("getMe", {})
        return _ok(result)
    except Exception as e:
        return _err(e)


# ─── Entry Point ──────────────────────────────────────────────────────────────

def main():
    """Main entry point for the tele-bot-mcp package script.

    Runs the MCP server over stdio, which is the standard transport
    used by MCP clients (Claude Desktop, agent orchestrators, etc.).
    """
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
