# app.py
import asyncio, logging
from contextlib import asynccontextmanager
import gradio as gr

from config import settings
from rabbit_base import RabbitBase
from listener import RabbitListenerBase
from rabbit_repo import RabbitRepo
from oa_server import OpenAIServers
#from vllm_backend import VLLMChatBackend, StubImagesBackend
#from transformers_backend import TransformersChatBackend, StubImagesBackend
#from hf_backend import HFChatBackend, StubImagesBackend
from timesfm_backend import TimesFMBackend, StubImagesBackend


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
log = logging.getLogger("app")

# ----------------- Hugging Face Spaces helpers -----------------
try:
    import spaces

    @spaces.GPU(duration=60)
    def gpu_entrypoint() -> str:
        return "gpu: ready"

except Exception:
    def gpu_entrypoint() -> str:
        return "gpu: not available (CPU only)"

# ----------------- RabbitMQ wiring -----------------
publisher = RabbitRepo(external_source="openai.mq.server")
resolver = (lambda name: "direct" if name.startswith("oa.") else settings.RABBIT_EXCHANGE_TYPE)
base = RabbitBase(exchange_type_resolver=resolver)

servers = OpenAIServers(
    publisher,
    chat_backend=TimesFMBackend(),
    images_backend=StubImagesBackend()
)

handlers = {
    "oaChatCreate": servers.handle_chat_create,
    "oaImagesGenerate": servers.handle_images_generate,
}

DECLS = [
    {"ExchangeName": "oa.chat.create", "FuncName": "oaChatCreate",
     "MessageTimeout": 600_000, "RoutingKeys": [settings.RABBIT_ROUTING_KEY]},
    {"ExchangeName": "oa.images.generate", "FuncName": "oaImagesGenerate",
     "MessageTimeout": 600_000, "RoutingKeys": [settings.RABBIT_ROUTING_KEY]},
]

listener = RabbitListenerBase(base, instance_name=settings.RABBIT_INSTANCE_NAME, handlers=handlers)

# ----------------- Service lifecycle -----------------
_startup_lock = asyncio.Lock()
_service_status = "Starting OpenAI MQ + TimesFM"
_listener_started = False


async def _start_listener() -> None:
    """Connect and register consumers once for this process."""
    global _listener_started, _service_status

    async with _startup_lock:
        if _listener_started and base.is_connected():
            return

        _service_status = "Connecting OpenAI MQ + TimesFM"
        try:
            await base.connect()
            await listener.start(DECLS)
        except Exception as exc:
            _listener_started = False
            _service_status = f"ERROR: {exc}"
            # A partial declaration must not leave a stale connection around if
            # the hosting platform retries this process.
            await base.close()
            log.exception("RabbitMQ listener startup failed")
            raise

        _listener_started = True
        _service_status = "OpenAI MQ + TimesFM: ready"
        log.info(_service_status)


async def _on_startup() -> None:
    # Fail the application startup if RabbitMQ is unavailable. Serving the UI
    # without its consumer makes a healthy-looking but non-functional service.
    await _start_listener()


async def _on_shutdown() -> None:
    global _listener_started, _service_status
    _listener_started = False
    _service_status = "Stopped"
    await base.close()


@asynccontextmanager
async def service_lifespan(app):
    # Pass this to launch: Gradio replaces demo.app and composes this lifespan
    # with its own lifecycle (including MCP) on the application it actually serves.
    await _on_startup()
    try:
        yield
    finally:
        await _on_shutdown()


async def service_status() -> str:
    """Read-only UI status; lifecycle startup is handled by FastAPI."""
    return _service_status


async def ping():
    return "ok" if _listener_started and base.is_connected() else _service_status

# ----------------- Gradio UI -----------------
with gr.Blocks(title="OpenAI over RabbitMQ (local vLLM)", theme=gr.themes.Soft()) as demo:
    gr.Markdown("## OpenAI-compatible over RabbitMQ — using vLLM locally inside Space")
    with gr.Tabs():
        with gr.Tab("Service"):
            btn = gr.Button("Ping")
            out = gr.Textbox(label="Ping result")
            btn.click(ping, inputs=None, outputs=out)
            init_status = gr.Textbox(label="Startup status", interactive=False)
            demo.load(fn=service_status, inputs=None, outputs=init_status)

        with gr.Tab("@spaces.GPU Probe"):
            gpu_btn = gr.Button("GPU Ready Probe", variant="primary")
            gpu_out = gr.Textbox(label="GPU Probe Result", interactive=False)
            gpu_btn.click(gpu_entrypoint, inputs=None, outputs=gpu_out)


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0", server_port=7860, show_error=True,
        debug=True, mcp_server=True,
        app_kwargs={"lifespan": service_lifespan},
    )