import argparse
import asyncio
from collections import defaultdict
import dataclasses
from functools import partial
import logging
import json
import os
import time
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union
import threading
import queue
import uuid

import httpx
import requests
from tenacity import retry, stop_after_attempt
from langport.core.base_node import BaseNode

from langport.protocol.worker_protocol import (
    GetNodeStateRequest,
    GetNodeStateResponse,
    NodeInfo,
    NodeInfoRequest,
    NodeInfoResponse,
    NodeListRequest,
    NodeListResponse,
    RegisterNodeRequest,
    RegisterNodeResponse,
    RemoveNodeRequest,
    RemoveNodeResponse,
    HeartbeatPing,
    HeartbeatPong,
)

from langport.constants import (
    HEART_BEAT_EXPIRATION,
    WORKER_API_TIMEOUT,
    WORKER_HEART_BEAT_INTERVAL,
    ErrorCode,
)
from langport.utils.cache_state import CacheState
from langport.utils.http_pool import AsyncHttpPool
from langport.utils.interval_timer import IntervalTimer
from cachetools import LRUCache, TTLCache
from asyncache import cached

class ClusterNode(BaseNode):
    def __init__(
        self,
        node_addr: str,
        node_id: str,
        init_neighborhoods_addr: List[str],
        logger: logging.Logger,
    ):
        super(ClusterNode, self).__init__(
            node_addr=node_addr,
            node_id=node_id,
            logger=logger,
        )
        self.init_neighborhoods_addr: List[str] = init_neighborhoods_addr
        self.neighborhoods: Dict[str, NodeInfo] = {}
        self.headers = {"User-Agent": "LangPort nodes"}

        self.states: Dict[str, CacheState] = {}
        self.remote_states: Dict[str, Dict[str, CacheState]] = defaultdict(dict)

        self.client = AsyncHttpPool()
        
        # timers
        self.add_timer(
            "heartbeat",
            WORKER_HEART_BEAT_INTERVAL,
            self.send_heartbeat_broadcast,
            args=None,
            kwargs=None,
            workers=1,
        )

        self.add_timer(
            "expiration_check",
            HEART_BEAT_EXPIRATION // 2,
            self.remove_nodes_expiration,
            args=None,
            kwargs=None,
            workers=1,
        )

        # start and stop
        self.on_start("get_all_init_neighborhoods", self.get_all_init_neighborhoods)
        self.on_start("register_node_broadcast", self.register_self_node_broadcast)
        self.on_stop("remove_node_broadcast", self.remove_self_node_broadcast)
    
    def _add_node(self, node_id: str, node_addr: str, check_heart_beat: bool=True):
        self.neighborhoods[node_id] = NodeInfo(
            node_id=node_id,
            node_addr=node_addr,
            check_heart_beat=check_heart_beat,
            refresh_time=int(time.time())
        )
    
    def _update_node(self, node_id: str, node_addr: str, check_heart_beat: bool=True):
        self.neighborhoods[node_id] = NodeInfo(
            node_id=node_id,
            node_addr=node_addr,
            check_heart_beat=check_heart_beat,
            refresh_time=int(time.time())
        )
    
    def _remove_node(self, node_id: str):
        if node_id in self.neighborhoods:
            del self.neighborhoods[node_id]
    
    async def get_all_init_neighborhoods(self):
        for neighbor_addr in self.init_neighborhoods_addr:
            try:
                info = await self.get_node_info(neighbor_addr)
            except httpx.ReadTimeout:
                self.logger.error(f"Cannot connect to neighbor {neighbor_addr}.")
                continue
            self._add_node(info.node_id, info.node_addr, info.check_heart_beat)
        # add self
        self._add_node(self.node_id, self.node_addr, True)
    
    async def get_node_info(self, node_addr: str) -> NodeInfo:
        response = await self.client.post(
            node_addr + "/node_info",
            headers=self.headers,
            json=NodeInfoRequest(node_id=self.node_id).dict(),
            timeout=WORKER_API_TIMEOUT,
        )
        remote_node_info = NodeInfoResponse.parse_obj(response.json())
        return remote_node_info.node_info

    async def register_node(self, target_node_addr: str, register_node_id: str, register_node_addr: str) -> bool:
        self.logger.info(f"Register {register_node_addr} to node {target_node_addr}")

        if target_node_addr == self.node_addr:
            self._add_node(register_node_id, register_node_addr)
            return True
        else:
            data = RegisterNodeRequest(
                node_id=register_node_id,
                node_addr=register_node_addr,
                check_heart_beat=True,
            )
        
            response = await self.client.post(
                target_node_addr + "/register_node",
                headers=self.headers,
                json=data.dict(),
                timeout=WORKER_API_TIMEOUT,
            )
            remote = RegisterNodeResponse.parse_obj(response.json())
            self._add_node(remote.node_id, remote.node_addr)
            
            # fetch remote node info
            remote_nodes = await self.fetch_all_nodes(remote.node_addr)
            for node in remote_nodes.nodes:
                self._update_node(node_id=node.node_id, node_addr=node.node_addr)
            return True

    async def register_node_broadcast(self, register_node_id: str, register_node_addr: str):
        self.logger.info(f"Register node broadcast. node_id: {register_node_id}, node_addr: {register_node_addr}")

        neighborhoods = [(k, v) for k, v in self.neighborhoods.items()]
        for node_id, node_info in neighborhoods:
            if node_info.node_addr == register_node_addr:
                continue
            await self.register_node(node_info.node_addr, register_node_id, register_node_addr)
    
    async def register_self_node_broadcast(self):
        await self.register_node_broadcast(self.node_id, self.node_addr)

    async def remove_node(self, target_node_addr: str, removed_node_id: str) -> bool:
        self.logger.info(f"Remove node {removed_node_id} from {target_node_addr}")

        if target_node_addr == self.node_addr:
            return True

        data = RemoveNodeRequest(
            node_id=removed_node_id,
        )

        response = await self.client.post(
            target_node_addr + "/remove_node",
            headers=self.headers,
            json=data.dict(),
            timeout=WORKER_API_TIMEOUT,
        )
        ret = RemoveNodeResponse.parse_obj(response.json())

        return True

    async def remove_node_broadcast(self, removed_node_id: str):
        self.logger.info(f"Remove node broadcast. node_addr: {removed_node_id}")
        neighborhoods = [(k, v) for k, v in self.neighborhoods.items()]
        for node_id, node_info in neighborhoods:
            if node_id == removed_node_id:
                continue
            await self.remove_node(node_info.node_addr, removed_node_id)

    async def remove_self_node_broadcast(self):
        await self.remove_node_broadcast(self.node_id)

    async def send_heartbeat(self, node_addr: str):
        data = HeartbeatPing(
            node_id=self.node_id,
        )

        response = await self.client.post(
            node_addr + "/heartbeat",
            headers=self.headers,
            json=data.dict(),
            timeout=WORKER_API_TIMEOUT,
        )
        ret = HeartbeatPong.parse_obj(response.json())
  
        return ret

    async def send_heartbeat_broadcast(self):
        self.logger.info(
            f"Send heartbeat. Worker: {self.node_id}; Address: {self.node_addr}."
        )

        self.logger.info(
            f"Neighborhoods: {[i for i, v in self.neighborhoods.items()]}."
        )

        for node_id, node_info in self.neighborhoods.items():
            if node_id == self.node_id:
                continue
            await self.send_heartbeat(node_info.node_addr)
        
    async def fetch_all_nodes(self, node_addr: str):
        data = NodeListRequest(
            node_id=self.node_id,
        )

        response = await self.client.post(
            node_addr + "/node_list",
            headers=self.headers,
            json=data.dict(),
            timeout=WORKER_API_TIMEOUT,
        )
        ret = NodeListResponse.parse_obj(response.json())
  
        return ret

    def remove_nodes_expiration(self):
        neighborhoods = [(k, v) for k, v in self.neighborhoods.items()]
        for node_id, node_info in neighborhoods:
            if node_id == self.node_id:
                continue
            if time.time() - node_info.refresh_time > HEART_BEAT_EXPIRATION:
                self.logger.info(f"Node {node_id} is expired.")
                self._remove_node(node_id)

    @cached(TTLCache(maxsize=1024, ttl=16))
    async def request_node_state(self, node_addr: str, name: str) -> GetNodeStateResponse:
        data = GetNodeStateRequest(
            state_name=name,
        )

        response = await self.client.post(
            node_addr + "/get_node_state",
            headers=self.headers,
            json=data.dict(),
            timeout=WORKER_API_TIMEOUT,
        )
        ret = GetNodeStateResponse.parse_obj(response.json())

        return ret
    
    async def _get_node_state_local(self, node_id: str, name: str) -> Any:
        entry = self.states.get(name, None)
        if entry is not None:
            return entry.get()
        else:
            return None
    
    @cached(TTLCache(maxsize=1024, ttl=1))
    async def _get_node_state_remote(self, node_id: str, name: str) -> Any:
        target_state = self.remote_states[node_id]
        if name in target_state:
            if target_state[name].is_valid():
                return target_state[name].get()

        node_info = self.neighborhoods[node_id]
        response = await self.request_node_state(node_info.node_addr, name)
        value = json.loads(response.state_value)
        ttl = response.state_ttl
        if value is not None:
            # print(f"refresh state {name}, node id {node_id}, value {value}")
            self.remote_states[node_id][name] = CacheState(value, ttl)
        return self.remote_states[node_id][name].get()
    
    async def get_node_state(self, node_id: str, name: str) -> Any:
        if node_id == self.node_id:
            return await self._get_node_state_local(node_id, name)
        else:
            state_lock = asyncio.Lock()
            async with state_lock:
                ret = await self._get_node_state_remote(node_id, name)
            return ret
    
    async def set_local_state(self, name: str, value: Any, ttl: int=60):
        self.states[name] = CacheState(value, ttl)

    async def api_register_node(self, request: RegisterNodeRequest) -> RegisterNodeResponse:
        if request.node_id not in self.neighborhoods:
            self.logger.info(f"Register {request.node_id} on {self.node_id}")
            self._add_node(request.node_id, request.node_addr)

            # broacast again
            await self.register_node_broadcast(request.node_id, request.node_addr)
        
        return RegisterNodeResponse(
            node_id=self.node_id,
            node_addr=self.node_addr,
            check_heart_beat=True
        )

    async def api_remove_node(self, request: RemoveNodeRequest) -> RemoveNodeResponse:
        if request.node_id in self.neighborhoods:
            self.logger.info(f"Remove {request.node_id} from {self.node_id}")
            self._remove_node(request.node_id)

            # broacast again
            await self.remove_node_broadcast(request.node_id)
        return RemoveNodeResponse(node_id=self.node_id)
    
    async def api_receive_heartbeat(self, request: HeartbeatPing) -> HeartbeatPong:
        if request.node_id in self.neighborhoods:
            self.neighborhoods[request.node_id].refresh_time = int(time.time())
        else:
            self.logger.info(f"Invalid ping packet from {request.node_id}.")
        return HeartbeatPong(exist=True)
    
    async def api_return_node_list(self, request: NodeListRequest) -> NodeListResponse:
        node_list = [node_info for node_id, node_info in self.neighborhoods.items()]
        return NodeListResponse(nodes=node_list)
    
    async def api_return_node_info(self, request: NodeInfoRequest) -> NodeInfoResponse:
        return NodeInfoResponse(
            node_info=NodeInfo(
                node_id=self.node_id,
                node_addr=self.node_addr,
                check_heart_beat=True,
            )
        )
    
    async def api_return_node_state(self, request: GetNodeStateRequest) -> GetNodeStateResponse:
        if request.state_name in self.states:
            entry = self.states[request.state_name]
            return GetNodeStateResponse(
                state_value=json.dumps(entry.value),
                state_ttl=entry.ttl
            )
        else:
            return GetNodeStateResponse(
                state_value=json.dumps(None)
            )