from __future__ import annotations

from typing import Optional
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from kilostar.utils.access import Accessor, TokenData
from kilostar.utils.ray_hook import ray_actor_hook

plugin_router = APIRouter(prefix="/api/v1/plugin", tags=["plugin"])


class SubmitRequest(BaseModel):
    org_name: str
    task_description: str
    context: Optional[dict] = None


@plugin_router.post("/submit")
async def submit_task(
    req: SubmitRequest,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    plugins = await pm.list_plugins.remote()
    if req.org_name not in plugins:
        raise HTTPException(404, f"Plugin '{req.org_name}' not found")

    org = ray_actor_hook(f"org_{req.org_name}").get(f"org_{req.org_name}")
    ctx = req.context or {}
    ctx["user"] = token_data.username
    task_id = await org.submit.remote(req.task_description, ctx)
    return {"task_id": task_id}


@plugin_router.get("/task/{task_id}")
async def get_task_status(
    task_id: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    db = ray_actor_hook("postgres_database").postgres_database
    task = await db.get_org_task.remote(task_id)
    if not task:
        raise HTTPException(404, "Task not found")
    return task


@plugin_router.get("/task/{task_id}/events")
async def get_task_events(
    task_id: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    db = ray_actor_hook("postgres_database").postgres_database
    events = await db.query_org_events.remote(task_id)
    return {"events": events}


@plugin_router.get("/task/{task_id}/stream")
async def stream_task(
    task_id: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    import asyncio

    org_name = None
    db = ray_actor_hook("postgres_database").postgres_database
    task = await db.get_org_task.remote(task_id)
    if not task:
        raise HTTPException(404, "Task not found")
    org_name = task["org_name"]

    org = ray_actor_hook(f"org_{org_name}").get(f"org_{org_name}")

    async def _generate():
        async for event in await org.stream.remote(task_id):
            yield f"data: {event}\n\n"

    return StreamingResponse(_generate(), media_type="text/event-stream")


@plugin_router.get("/list")
async def list_plugins(
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    plugins = await pm.list_plugins.remote()
    return {"plugins": plugins}


@plugin_router.post("/install")
async def install_plugin(
    name: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    await pm.install.remote(name)
    return {"status": "ok", "name": name}


@plugin_router.post("/reload/{name}")
async def reload_plugin(
    name: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    await pm.reload.remote(name)
    return {"status": "ok", "name": name}