from __future__ import annotations
import uuid
import asyncio
import json
from loguru import logger
from service_forge.workflow.trigger import Trigger
from service_forge.workflow.trigger_event import TriggerEvent
from typing import AsyncIterator, Any
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState
from service_forge.workflow.port import Port
from google.protobuf.message import Message
from google.protobuf.json_format import MessageToJson

class WebSocketAPIV2Trigger(Trigger):
    DEFAULT_INPUT_PORTS = [
        Port("app", FastAPI),
        Port("path", str),
    ]

    DEFAULT_OUTPUT_PORTS = [
        Port("websocket", WebSocket),
        Port("user_id", str),
        Port("token", str),
    ]

    def __init__(self, name: str):
        super().__init__(name)
        self.events = {}
        self.is_setup_websocket = False

    def _setup_websocket(self, app: FastAPI, path: str) -> None:
        async def websocket_handler(websocket: WebSocket):
            websocket.state.user_id = websocket.headers.get("X-User-ID") or "0"
            websocket.state.auth_token = websocket.headers.get("X-User-Token") or ""

            logger.info(f'user_id {websocket.state.user_id} token {websocket.state.auth_token}')
            
            await websocket.accept()

            task_id = uuid.uuid4()
            self.result_queues[task_id] = asyncio.Queue()

            self.trigger_queue.put_nowait({
                "id": task_id,
                "websocket": websocket,
                "user_id": websocket.state.user_id,
                "token": websocket.state.auth_token,
                "data": {},
            })

            await self.result_queues[task_id].get()

        app.websocket(path)(websocket_handler)

    async def _run(self, app: FastAPI, path: str) -> AsyncIterator[bool]:
        if not self.is_setup_websocket:
            self._setup_websocket(app, path)
            self.is_setup_websocket = True

        while True:
            try:
                trigger = await self.trigger_queue.get()
                self.prepare_output_edges('user_id', trigger['user_id'])
                self.prepare_output_edges('token', trigger['token'])
                self.prepare_output_edges('websocket', trigger['websocket'])
                event = TriggerEvent(
                    trigger_node=self,
                    task_id=trigger['id'],
                    trace_context=None,
                    user_id=trigger['user_id'],
                    data=trigger,
                )
                self.trigger(event)
                yield event
            except Exception as e:
                logger.error(f"Error in WebSocketAPITrigger._run: {e}")
                continue

    async def _stop(self) -> AsyncIterator[bool]:
        pass