#!/usr/bin/env python3
# Zabbix Alert Script: Rich Telegram Notification
# Sends HTML formatted alert + graph image
# Input: JSON String from Zabbix
# Place in: /usr/lib/zabbix/alertscripts/

import sys
import requests
import json
import logging

# --- CONFIGURATION ---
ZABBIX_URL = "http://localhost"  # Base URL, no trailing slash
ZABBIX_USER = "Admin"
ZABBIX_PASS = "zabbix"
TG_BOT_TOKEN = "YOUR_BOT_TOKEN"
LOG_FILE = "/var/log/zabbix/telegram_bot.log"
# ---------------------

# Setup logging
logging.basicConfig(filename=LOG_FILE, level=logging.INFO, 
                    format='%(asctime)s [%(levelname)s] %(message)s')

def get_emoji(severity):
    # Map Zabbix severity (0-5) to Emojis
    # 0: Not classified, 1: Info, 2: Warning, 3: Average, 4: High, 5: Disaster
    # Also handles "Recovery" status logic in message builder
    mapping = {
        '0': '📓', # Not classified
        '1': 'ℹ️', # Information
        '2': '⚠️', # Warning
        '3': '🟧', # Average
        '4': '🔥', # High
        '5': '🛑'  # Disaster
    }
    return mapping.get(str(severity), '❓')

def login_zabbix():
    api_url = f"{ZABBIX_URL}/api_jsonrpc.php"
    headers = {'Content-Type': 'application/json'}
    data = {
        "jsonrpc": "2.0",
        "method": "user.login",
        "params": {
            "username": ZABBIX_USER,
            "password": ZABBIX_PASS
        },
        "id": 1
    }
    
    try:
        resp = requests.post(api_url, json=data, timeout=10)
        result = resp.json().get('result')
        if not result:
            logging.error(f"Zabbix Login Failed: {resp.text}")
        return result
    except Exception as e:
        logging.error(f"Zabbix Login Exception: {e}")
        return None

def get_graph_image(session_id, item_id, period=3600):
    if not item_id:
        return None
        
    chart_url = f"{ZABBIX_URL}/chart3.php"
    cookies = {'zbx_session': session_id}
    params = {
        'items[0][itemid]': item_id,
        'items[0][drawtype]': 5, # Gradient
        'items[0][color]': 'FF5555',
        'period': period,
        'width': 600,
        'height': 200,
        'legend': 1
    }
    
    try:
        resp = requests.get(chart_url, params=params, cookies=cookies, timeout=15)
        if resp.status_code == 200 and 'image' in resp.headers.get('Content-Type', ''):
            return resp.content
        logging.warning(f"Failed to fetch graph. Status: {resp.status_code}")
    except Exception as e:
        logging.error(f"Graph Fetch Exception: {e}")
    return None

def send_telegram(chat_id, message_html, image_data=None):
    base_url = f"https://api.telegram.org/bot{TG_BOT_TOKEN}"
    
    try:
        # 1. Send Text
        # Note: If image is present, usually nicer to send image with caption, 
        # but captions have length limits (1024 chars). 
        # Safest is to send text first, then image.
        
        # However, for a cleaner look ("Option A"), we want text + image.
        # Let's try sending Photo with Caption if text is short enough, 
        # otherwise split.
        
        if image_data and len(message_html) < 1000:
            url_photo = f"{base_url}/sendPhoto"
            files = {'photo': ('graph.png', image_data, 'image/png')}
            data = {'chat_id': chat_id, 'caption': message_html, 'parse_mode': 'HTML'}
            resp = requests.post(url_photo, data=data, files=files, timeout=15)
        else:
            # Send Message
            url_msg = f"{base_url}/sendMessage"
            data = {'chat_id': chat_id, 'text': message_html, 'parse_mode': 'HTML'}
            resp = requests.post(url_msg, data=data, timeout=15)
            
            # Send Image separately if exists
            if image_data:
                url_photo = f"{base_url}/sendPhoto"
                files = {'photo': ('graph.png', image_data, 'image/png')}
                requests.post(url_photo, data={'chat_id': chat_id}, files=files, timeout=15)
                
        if resp.status_code != 200:
            logging.error(f"Telegram Send Failed: {resp.text}")
            
    except Exception as e:
        logging.error(f"Telegram Exception: {e}")

def format_message(data):
    # Extract data with safe defaults
    status = data.get('event_status', 'PROBLEM') # PROBLEM or RESOLVED
    severity = data.get('event_nseverity', '0')
    event_name = data.get('event_name', 'Unknown Event')
    host_name = data.get('host_name', 'Unknown Host')
    host_ip = data.get('host_ip', '')
    event_time = f"{data.get('event_date', '')} {data.get('event_time', '')}"
    item_val = data.get('item_value', 'N/A')
    event_id = data.get('event_id', '')
    trigger_id = data.get('trigger_id', '')
    opdata = data.get('event_opdata', '')

    emoji = get_emoji(severity)
    
    # Template V2 - Action Oriented
    if status == 'RESOLVED':
        header = f"✅ <b>RESOLVED: {event_name}</b>"
        # More clean/concise for resolved
        msg = f"{header}\n\n"
        msg += f"📍 <b>Host:</b> {host_name}\n"
        msg += f"⏱️ <b>Time:</b> {event_time}\n"
    else:
        header = f"{emoji} <b>{event_name}</b>"
        
        msg = f"{header}\n\n"
        msg += f"📍 <b>Host:</b> {host_name} ({host_ip})\n"
        
        # Prefer OpData if available (shows full context like "CPU: 99% > 90%")
        # Otherwise show raw Value
        if opdata:
            msg += f"📊 <b>Dados:</b> {opdata}\n"
        else:
            msg += f"📉 <b>Valor:</b> <code>{item_val}</code>\n"
            
        msg += f"\n<b>🛠️ AÇÕES RÁPIDAS:</b>\n"
        
        # Action Links
        if event_id and trigger_id:
            link_event = f"{ZABBIX_URL}/tr_events.php?triggerid={trigger_id}&eventid={event_id}"
            link_ack = f"{ZABBIX_URL}/zabbix.php?action=acknowledge.edit&eventids[]={event_id}"
            
            msg += f"🔴 <a href=\"{link_event}\">Ver Evento no Zabbix</a>\n"
            msg += f"🛡️ <a href=\"{link_ack}\">Reconhecer (Ack)</a>"
            
    return msg

if __name__ == "__main__":
    # Zabbix sends parameters:
    # $1 = SendTo (Chat ID)
    # $2 = Subject (Ignored in this version, we build our own)
    # $3 = Message (Expects JSON String)
    
    if len(sys.argv) < 4:
        logging.error("Missing arguments. Usage: script.py <chat_id> <subject> <json_data>")
        print("Usage: telegram_graph.py <chat_id> <subject> <json_data>")
        sys.exit(1)

    chat_id = sys.argv[1]
    # subject = sys.argv[2] # Unused
    raw_data = sys.argv[3]

    try:
        data = json.loads(raw_data)
    except json.JSONDecodeError:
        logging.error(f"Invalid JSON received: {raw_data}")
        # Fallback: Send raw text if JSON fails
        send_telegram(chat_id, f"⚠️ <b>JSON Error</b>\n{raw_data}")
        sys.exit(1)

    # Format HTML
    message_html = format_message(data)

    # Fetch Graph
    graph_img = None
    item_id = data.get('item_id')
    
    # Only fetch graph for "PROBLEM"
    if data.get('event_status') != 'RESOLVED' and item_id:
        token = login_zabbix()
        if token:
            graph_img = get_graph_image(token, item_id)
    
    # Send
    send_telegram(chat_id, message_html, graph_img)