from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from fastapi import Query

from pymongo.asynchronous.database import AsyncDatabase
from extra.NodeDTO import NodeDTO
from typing import List, Annotated

class WebExtension:
    app: FastAPI
    dbStore: AsyncDatabase
    def __init__(self, core):
        self.core = core
        self.app = core.app
        self.dbStore = core.dbStore
    
        @self.app.get(f"{self.core.context}/nodes/list")
        async def listOfNodes():
            pipeline = [ #ai slooop
                # Сортировка для каждого num по убыванию ts
                {"$sort": {"num": 1, "ts": -1}},
                # Группировка по num и взятие первого (последнего) документа
                {"$group": {
                    "_id": "$num",
                    "latest_doc": {"$first": "$$ROOT"}
                }},
                # Замена корневого документа
                {"$replaceRoot": {"newRoot": "$latest_doc"}},
                # Сортировка результата (опционально)
                {"$sort": {"num": 1}}
            ]

            collection = self.dbStore['node_info']
            c = await collection.aggregate(pipeline)
            l = await c.to_list()
            return [NodeDTO(node) for node in l]
        
        @self.app.get(f"{self.core.context}/nodes/direct")
        async def listOfDirectNodes():
            pipeline = [ #ai slooop
                {"$match": {"hops_away": 0}},  # Фильтруем по списку
                # Сортировка для каждого num по убыванию ts
                {"$sort": {"num": 1, "ts": -1}},
                # Группировка по num и взятие первого (последнего) документа
                {"$group": {
                    "_id": "$num",
                    "latest_doc": {"$first": "$$ROOT"}
                }},
                # Замена корневого документа
                {"$replaceRoot": {"newRoot": "$latest_doc"}},
                # Сортировка результата (опционально)
                {"$sort": {"num": 1}}
            ]

            collection = self.dbStore['node_info']
            c = await collection.aggregate(pipeline)
            l = await c.to_list()
            return [NodeDTO(node) for node in l]
        
        @self.app.get(self.core.context + "/nodes/{num}")
        async def oneNode(num: int):
            collection = self.dbStore['node_info']
            c = await collection.find_one(
                {"num":num},
                sort=[("ts", -1)]
            )
            if c:
                return NodeDTO(c)
            else:
                raise HTMLResponse(status_code=404)
            
        @self.app.get(self.core.context + "/nodes")
        async def listOfSelectedNodes(nums: List[int] = Query(None)):
            pipeline = [
                {"$match": {"num": {"$in": nums}}},  # Фильтруем по списку
                {"$sort": {"num": 1, "ts": -1}},     # Сортируем по num и ts (по убыванию)
                {
                    "$group": {
                        "_id": "$num",                # Группируем по num
                        "latest_doc": {"$first": "$$ROOT"}  # Берем первую (последнюю по ts)
                    }
                },
                {"$replaceRoot": {"newRoot": "$latest_doc"}},  # Восстанавливаем структуру
                {"$sort": {"num": 1}}  # Сортируем результат по num
            ]
            collection = self.dbStore['node_info']
            c = await collection.aggregate(pipeline)
            l = await c.to_list()
            return [NodeDTO(node) for node in l]