File size: 14,992 Bytes
915d8f9
 
0f4be73
3222a21
915d8f9
bfe88a9
3222a21
915d8f9
 
bfe88a9
915d8f9
 
 
 
 
527a73c
915d8f9
 
bfe88a9
 
 
 
915d8f9
bfe88a9
915d8f9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3222a21
915d8f9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3222a21
915d8f9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
# app/main.py
import gradio as gr
import httpx
import websockets
import asyncio
import json
import os
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Depends # Import FastAPI itself
from .database import connect_db, disconnect_db, database, metadata, users
from .api import router as api_router
from . import schemas, auth, dependencies
from .websocket import manager # Import the connection manager instance

from sqlalchemy.schema import CreateTable
from sqlalchemy.dialects import sqlite

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

API_BASE_URL = "http://127.0.0.1:7860/api"

# --- Lifespan (remains the same) ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    # ... (same DB setup code) ...
    logger.info("Application startup: Connecting DB...")
    await connect_db()
    logger.info("Application startup: DB Connected. Checking/Creating tables...")
    if database.is_connected:
        try:
            check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;"
            table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name})
            if not table_exists:
                logger.info(f"Table '{users.name}' not found, attempting creation using async connection...")
                dialect = sqlite.dialect()
                create_table_stmt = str(CreateTable(users).compile(dialect=dialect))
                await database.execute(query=create_table_stmt)
                logger.info(f"Table '{users.name}' created successfully via async connection.")
                table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name})
                if table_exists_after: logger.info(f"Table '{users.name}' verified after creation.")
                else: logger.error(f"Table '{users.name}' verification FAILED after creation attempt!")
            else:
                logger.info(f"Table '{users.name}' already exists (checked via async connection).")
        except Exception as db_setup_err:
            logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}")
    else:
        logger.error("CRITICAL: Database connection failed, skipping table setup.")
    logger.info("Application startup: DB setup phase complete.")
    yield
    logger.info("Application shutdown: Disconnecting DB...")
    await disconnect_db()
    logger.info("Application shutdown: DB Disconnected.")


# --- FastAPI App Setup (remains the same) ---
app = FastAPI(lifespan=lifespan)
app.include_router(api_router, prefix="/api")

# --- Helper functions (make_api_request remains the same) ---
async def make_api_request(method: str, endpoint: str, **kwargs):
    async with httpx.AsyncClient() as client:
        url = f"{API_BASE_URL}{endpoint}"
        try: response = await client.request(method, url, **kwargs); response.raise_for_status(); return response.json()
        except httpx.RequestError as e: logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}"); return {"error": f"Network error: {e}"}
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
            try: detail = e.response.json().get("detail", e.response.text)
            except json.JSONDecodeError: detail = e.response.text
            return {"error": f"API Error: {detail}"}
        except Exception as e: logger.error(f"Unexpected error during API call: {e}"); return {"error": f"Unexpected error: {str(e)}"}

# --- WebSocket handling ---
# <<< Pass state objects by reference >>>
async def listen_to_websockets(token: str, notification_list_state: gr.State, notification_trigger_state: gr.State):
    """Connects to WS and updates state list and trigger when a message arrives."""
    ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
    logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")

    if not token:
        logger.warning(f"[{ws_listener_id}] No token provided. Task exiting.")
        return # Just exit, don't need to return state values

    ws_url_base = API_BASE_URL.replace("http", "ws")
    ws_url = f"{ws_url_base}/ws/{token}"
    logger.info(f"[{ws_listener_id}] Attempting to connect: {ws_url}")

    try:
        async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
            logger.info(f"[{ws_listener_id}] WebSocket connected successfully.")
            while True:
                try:
                    message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0)
                    logger.info(f"[{ws_listener_id}] Received: {message_str[:100]}...")
                    try:
                         message_data = json.loads(message_str)
                         if message_data.get("type") == "new_user":
                            notification = schemas.Notification(**message_data)
                            logger.info(f"[{ws_listener_id}] Processing 'new_user': {notification.message}")
                            # --- Modify state values directly ---
                            current_list = notification_list_state.value.copy() # Operate on copy
                            current_list.insert(0, notification.message)
                            if len(current_list) > 10: current_list.pop()
                            notification_list_state.value = current_list # Assign back modified copy
                            notification_trigger_state.value += 1 # Increment trigger
                            # --- Log the update ---
                            logger.info(f"[{ws_listener_id}] States updated: list len={len(notification_list_state.value)}, trigger={notification_trigger_state.value}")
                         else:
                              logger.warning(f"[{ws_listener_id}] Unknown message type: {message_data.get('type')}")
                    # ... (error handling for parsing) ...
                    except json.JSONDecodeError: logger.error(f"[{ws_listener_id}] JSON Decode Error: {message_str}")
                    except Exception as parse_err: logger.error(f"[{ws_listener_id}] Message Processing Error: {parse_err}")
                # ... (error handling for websocket recv/connection) ...
                except asyncio.TimeoutError: logger.debug(f"[{ws_listener_id}] WebSocket recv timed out."); continue
                except websockets.ConnectionClosedOK: logger.info(f"[{ws_listener_id}] WebSocket closed normally."); break
                except websockets.ConnectionClosedError as e: logger.error(f"[{ws_listener_id}] WebSocket closed with error: {e}"); break
                except Exception as e: logger.error(f"[{ws_listener_id}] Listener loop error: {e}"); await asyncio.sleep(1); break
    # ... (error handling for websocket connect) ...
    except asyncio.TimeoutError: logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out.")
    except websockets.exceptions.InvalidURI: logger.error(f"[{ws_listener_id}] Invalid WebSocket URI.")
    except websockets.exceptions.WebSocketException as e: logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
    except Exception as e: logger.error(f"[{ws_listener_id}] Unexpected error in listener task: {e}")

    logger.info(f"[{ws_listener_id}] Listener task finished.")
    # No need to return state values when task ends

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # State variables
    auth_token = gr.State(None)
    user_info = gr.State(None)
    notification_list = gr.State([])
    websocket_task = gr.State(None)
    # Add trigger states
    notification_trigger = gr.State(0)
    last_polled_trigger = gr.State(0) # State to track last seen trigger

    # --- UI Components ---
    with gr.Tabs() as tabs:
        # --- Registration/Login Tabs (remain the same) ---
        with gr.TabItem("Register", id="register_tab"):
            gr.Markdown("## Create a new account")
            reg_email = gr.Textbox(label="Email", type="email")
            reg_password = gr.Textbox(label="Password (min 8 chars)", type="password")
            reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
            reg_button = gr.Button("Register")
            reg_status = gr.Textbox(label="Status", interactive=False)
        with gr.TabItem("Login", id="login_tab"):
            gr.Markdown("## Login to your account")
            login_email = gr.Textbox(label="Email", type="email")
            login_password = gr.Textbox(label="Password", type="password")
            login_button = gr.Button("Login")
            login_status = gr.Textbox(label="Status", interactive=False)

        # --- Welcome Tab ---
        with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
            gr.Markdown("## Welcome!", elem_id="welcome_header")
            welcome_message = gr.Markdown("", elem_id="welcome_message")
            logout_button = gr.Button("Logout")
            gr.Markdown("---")
            gr.Markdown("## Real-time Notifications")
            notification_display = gr.Textbox( # Visible display
                label="New User Alerts", lines=5, max_lines=10, interactive=False
            )
            # <<< Hidden component for polling >>>
            dummy_poller = gr.Number(label="poller", value=0, visible=False, every=1)


    # --- Event Handlers ---

    # Registration Logic (remains the same)
    async def handle_register(email, password, confirm_password):
        if not email or not password or not confirm_password: return gr.update(value="Please fill fields.")
        if password != confirm_password: return gr.update(value="Passwords mismatch.")
        if len(password) < 8: return gr.update(value="Password >= 8 chars.")
        payload = {"email": email, "password": password}
        result = await make_api_request("post", "/register", json=payload)
        if "error" in result: return gr.update(value=f"Register failed: {result['error']}")
        else: return gr.update(value=f"Register success: {result.get('email')}! Log in.")
    reg_button.click(handle_register, inputs=[reg_email, reg_password, reg_confirm_password], outputs=[reg_status])

    # Login Logic
    # <<< MODIFIED: Pass/Reset both trigger states >>>
    async def handle_login(email, password, current_task, current_trigger_val, current_last_poll_val):
        # Define failure outputs matching the outputs list
        fail_outputs = (gr.update(value="..."), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val, current_last_poll_val)

        if not email or not password: return fail_outputs[:1] + (gr.update(value="Enter email/password"),) + fail_outputs[2:]

        payload = {"email": email, "password": password}
        result = await make_api_request("post", "/login", json=payload)

        if "error" in result: return (gr.update(value=f"Login failed: {result['error']}"),) + fail_outputs[1:]
        else:
            token = result.get("access_token")
            user_data = await dependencies.get_optional_current_user(token)
            if not user_data: return (gr.update(value="Login ok but user fetch failed."),) + fail_outputs[1:]

            if current_task and not current_task.done():
                current_task.cancel()
                try: await current_task
                except asyncio.CancelledError: logger.info("Previous WebSocket task cancelled.")

            # <<< Pass state objects to listener >>>
            new_task = asyncio.create_task(listen_to_websockets(token, notification_list, notification_trigger))

            welcome_msg = f"Welcome, {user_data.email}!"
            # Reset triggers on successful login
            return (
                gr.update(value="Login successful!"), token, user_data.model_dump(), # status, token, user_info
                gr.update(selected="welcome_tab"), gr.update(visible=True), gr.update(value=welcome_msg), # UI changes
                new_task, 0, 0 # websocket_task, notification_trigger, last_polled_trigger
            )

    # <<< MODIFIED: Add last_polled_trigger to inputs/outputs >>>
    login_button.click(
        handle_login,
        inputs=[login_email, login_password, websocket_task, notification_trigger, last_polled_trigger],
        outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger, last_polled_trigger]
    )


    # <<< Polling function >>>
    def poll_and_update(current_trigger_value, last_known_trigger, current_notif_list):
        """ Checks trigger state change and updates UI if needed. """
        if current_trigger_value != last_known_trigger:
            logger.info(f"Polling detected trigger change ({last_known_trigger} -> {current_trigger_value}). Updating UI.")
            new_value = "\n".join(current_notif_list)
            # Return new display value AND update last_known_trigger state value
            return gr.update(value=new_value), current_trigger_value
        else:
            # No change, return NoUpdate for display AND existing last_known_trigger value
            return gr.NoUpdate(), last_known_trigger

    # <<< Attach polling function to the dummy component's change event >>>
    dummy_poller.change(
        fn=poll_and_update,
        inputs=[notification_trigger, last_polled_trigger, notification_list],
        outputs=[notification_display, last_polled_trigger], # Update display & last polled state
        queue=False # Try immediate update
    )


    # Logout Logic
    # <<< MODIFIED: Reset both trigger states >>>
    async def handle_logout(current_task):
        if current_task and not current_task.done():
            current_task.cancel()
            try: await current_task
            except asyncio.CancelledError: logger.info("WebSocket task cancelled on logout.")
        # Reset all relevant states
        return ( None, None, [], None, # auth_token, user_info, notification_list, websocket_task
                 gr.update(selected="login_tab"), gr.update(visible=False), # tabs, welcome_tab
                 gr.update(value=""), gr.update(value=""), # welcome_message, login_status
                 0, 0 ) # notification_trigger, last_polled_trigger (reset)

    # <<< MODIFIED: Add last_polled_trigger to outputs >>>
    logout_button.click(
        handle_logout,
        inputs=[websocket_task],
        outputs=[
            auth_token, user_info, notification_list, websocket_task,
            tabs, welcome_tab, welcome_message, login_status,
            notification_trigger, last_polled_trigger # Add trigger states here
        ]
    )

# Mount Gradio App (remains the same)
app = gr.mount_gradio_app(app, demo, path="/")

# Run Uvicorn (remains the same)
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=True)