import os
from telethon import TelegramClient
from telethon.errors import SessionPasswordNeededError
import networkx as nx
import matplotlib.pyplot as plt
import gradio as gr

# Obtener credenciales de las variables de entorno
API_ID = os.getenv("API_ID")
API_HASH = os.getenv("API_HASH")

print(f"API_ID: {API_ID}, API_HASH: {API_HASH}")  # Mensaje de depuración

if not API_ID or not API_HASH:
    raise ValueError("API_ID or API_HASH is missing. Please check your environment variables.")

try:
    API_ID = int(API_ID)  # Convertir API_ID a entero
except ValueError:
    raise ValueError("API_ID must be a valid integer.")

# Configurar el cliente de Telegram con una sesión persistente
client = TelegramClient('my_session', API_ID, API_HASH)

async def fetch_messages(channel_username):
    try:
        await client.start()
        messages = []
        async for message in client.iter_messages(channel_username, limit=50):
            messages.append({
                'id': message.id,
                'text': message.text[:100] + "..." if message.text else "No text",
                'date': str(message.date),
                'forwarded_from': message.forward.sender.username if message.forward else None
            })
        print(f"Messages fetched: {len(messages)}")  # Mensaje de depuración
        return messages
    except Exception as e:
        print(f"Error fetching messages: {str(e)}")
        return []  # Devuelve una lista vacía en caso de error

def build_network(messages):
    G = nx.DiGraph()
    for msg in messages:
        origin = msg['forwarded_from']
        destination = 'TargetChannel'  # Canal actual
        if origin and destination:
            if G.has_edge(origin, destination):
                G[origin][destination]['weight'] += 1
            else:
                G.add_edge(origin, destination, weight=1)
    return G

def visualize_network(channel_username):
    try:
        import asyncio
        messages = asyncio.run(fetch_messages(channel_username))
        
        if not messages:  # Si no hay mensajes, devuelve un gráfico vacío
            plt.figure(figsize=(10, 8))
            plt.title("No data available for this channel")
            plt.text(0.5, 0.5, "Try another channel or check the username.", fontsize=12, ha='center')
            return plt.gcf()
        
        G = build_network(messages)
        
        if G.number_of_edges() == 0:  # Si no hay aristas, muestra un mensaje
            plt.figure(figsize=(10, 8))
            plt.title("No forwarding patterns found")
            plt.text(0.5, 0.5, "This channel has no forwarded messages.", fontsize=12, ha='center')
            return plt.gcf()
        
        plt.figure(figsize=(10, 8))
        pos = nx.spring_layout(G)
        nx.draw(G, pos, with_labels=True, node_size=3000, node_color='lightblue', font_size=10, font_weight='bold')
        labels = nx.get_edge_attributes(G, 'weight')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=labels)
        plt.title("Telegram Channel Network")
        return plt.gcf()  # Devuelve la figura actual
    except Exception as e:
        print(f"Error in visualize_network: {str(e)}")
        plt.figure(figsize=(10, 8))
        plt.title("An error occurred")
        plt.text(0.5, 0.5, "Please check the logs for more details.", fontsize=12, ha='center')
        return plt.gcf()

# Interfaz de Gradio
demo = gr.Interface(
    fn=visualize_network,
    inputs="text",
    outputs="plot",
    title="Telegram Network Mapper",
    description="Enter a Telegram channel username to analyze its message forwarding network."
)

# Desactivar SSR y lanzar la aplicación
demo.launch(share=False, ssr_mode=False)