
from __future__ import annotations
from service_forge.workflow.trigger import Trigger
from typing import AsyncIterator, Any
from fastapi import FastAPI, Request
import uuid
import asyncio
import json
from service_forge.workflow.port import Port
from service_forge.utils.default_type_converter import type_converter
from service_forge.api.websocket_manager import websocket_manager
from fastapi import HTTPException
from google.protobuf.message import Message
from google.protobuf.json_format import MessageToJson

class FastAPITrigger(Trigger):
    DEFAULT_INPUT_PORTS = [
        Port("app", FastAPI),
        Port("path", str),
        Port("method", str),
        Port("data_type", type),
    ]

    DEFAULT_OUTPUT_PORTS = [
        Port("trigger", bool),
        Port("user_id", int),
        Port("data", Any),
    ]

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

    @staticmethod
    def serialize_result(result: Any):
        if isinstance(result, Message):
            return MessageToJson(
                result,
                preserving_proto_field_name=True
            )
        return result

    async def handle_request(
        self,
        request: Request,
        data_type: type,
        extract_data_fn: callable[[Request], dict]
    ):
        task_id = uuid.uuid4()
        self.result_queues[task_id] = asyncio.Queue()

        body_data = await extract_data_fn(request)
        converted_data = data_type(**body_data)

        client_id = (
            body_data.get("client_id")
            or request.query_params.get("client_id")
            or request.headers.get("X-Client-ID")
        )
        if client_id:
            workflow_name = getattr(self.workflow, "name", "Unknown")
            steps = len(self.workflow.nodes) if hasattr(self.workflow, "nodes") else 1
            websocket_manager.create_task_with_client(task_id, client_id, workflow_name, steps)

        self.trigger_queue.put_nowait({
            "id": task_id,
            "user_id": getattr(request.state, "user_id", None),
            "data": converted_data,
        })

        result = await self.result_queues[task_id].get()
        del self.result_queues[task_id]

        if isinstance(result, HTTPException):
            raise result
        if isinstance(result, Exception):
            raise HTTPException(status_code=500, detail=str(result))

        return self.serialize_result(result)

    def _setup_route(self, app: FastAPI, path: str, method: str, data_type: type) -> None:
        async def get_data(request: Request) -> dict:
            return dict(request.query_params)

        async def body_data(request: Request) -> dict:
            raw = await request.body()
            if not raw:
                return {}
            return json.loads(raw.decode("utf-8"))

        extractor = get_data if method == "GET" else body_data

        async def handler(request: Request):
            return await self.handle_request(request, data_type, extractor)

        if method == "GET":
            app.get(path)(handler)
        elif method == "POST":
            app.post(path)(handler)
        elif method == "PUT":
            app.put(path)(handler)
        elif method == "DELETE":
            app.delete(path)(handler)
        else:
            raise ValueError(f"Invalid method {method}")

    async def _run(self, app: FastAPI, path: str, method: str, data_type: type) -> AsyncIterator[bool]:
        if not self.is_setup_route:
            self._setup_route(app, path, method, data_type)
            self.is_setup_route = True

        while True:
            try:
                trigger = await self.trigger_queue.get()
                self.prepare_output_edges(self.get_output_port_by_name('user_id'), trigger['user_id'])
                self.prepare_output_edges(self.get_output_port_by_name('data'), trigger['data'])
                yield self.trigger(trigger['id'])
            except Exception as e:
                from loguru import logger
                logger.error(f"Error in FastAPITrigger._run: {e}")
                continue
