import asyncio
import os
from contextlib import asynccontextmanager

import uvicorn
from connections.rabbitmq import rabbitmq_manager
from fasthtml.core import FastHTML, Mount
from modules.fasthtml import app as fasthtml_app
from modules.rabbitmq import consume_messages_topic, publish_heartbeat


async def start_heartbeat():
    """Start heartbeat with auto-reconnect"""
    while True:
        try:
            await publish_heartbeat()
        except Exception as e:
            print(f"Heartbeat failed: {e}")
            await asyncio.sleep(5)


async def start_consumer():
    """Start consumer with auto-reconnect"""
    while True:
        try:
            await consume_messages_topic()
        except Exception as e:
            print(f"Consumer failed: {e}")
            await asyncio.sleep(5)


async def connect_with_retry():
    """Retry RabbitMQ connection until successful"""
    while True:
        try:
            await rabbitmq_manager.connect()
            if (
                rabbitmq_manager._connection
                and not rabbitmq_manager._connection.is_closed
            ):
                print("RabbitMQ connected")
                asyncio.create_task(start_heartbeat())
                asyncio.create_task(start_consumer())
                return
            else:
                await asyncio.sleep(5)
        except Exception as e:
            print(f"RabbitMQ connection failed: {e}, retrying in 5 seconds...")
            await asyncio.sleep(5)


@asynccontextmanager
async def lifespan(app):
    print("Starting up...")

    asyncio.create_task(connect_with_retry())

    yield

    print("Shutting down...")
    for task in asyncio.all_tasks():
        if task is not asyncio.current_task():
            task.cancel()

    await rabbitmq_manager.close()


app = FastHTML(routes=[Mount("", fasthtml_app, name="FastHTML")])
app.set_lifespan(lifespan)

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=int(os.environ["PORT"]), reload=True)