import argparse
import atexit
import json
import logging
import threading
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import requests
from fastapi import FastAPI, HTTPException, Request, status
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import Hf3fsMetadataInterface

# --- Configuration ---
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


# --- Data Models ---
class RankMetadata:
    """Holds all metadata for a single rank."""

    def __init__(self, num_pages: int):
        self.lock = threading.RLock()
        self.num_pages = num_pages
        self.free_pages: List[int] = list(range(num_pages))
        self.key_to_index: Dict[str, int] = {}
        # Todo: Support multi files for HF3FS

    def exists_keys(self, keys: List[str]) -> List[bool]:
        """Check if keys exist in metadata."""
        with self.lock:
            return [key in self.key_to_index for key in keys]

    def reserve_and_allocate_page_indices(
        self, keys: List[Tuple[str, str]]
    ) -> List[Tuple[bool, int]]:
        """Reserve and allocate page indices for keys."""
        with self.lock:
            results = [None] * len(keys)
            new_keys_to_process = []

            for i, (key, prefix_key) in enumerate(keys):
                if key in self.key_to_index:
                    results[i] = (True, self.key_to_index[key])
                else:
                    new_keys_to_process.append((i, key, prefix_key))

            # Todo: Implementing data eviction logic after HiCache supports prefix information pass-through
            for i, key, prefix_key in new_keys_to_process:
                if len(self.free_pages) > 0:
                    page_idx = self.free_pages.pop()
                    results[i] = (False, page_idx)
                else:
                    results[i] = (False, -1)

            return results

    def confirm_write(
        self,
        written_keys_to_confirm: List[Tuple[str, int]],
        pages_to_release: List[int],
    ) -> None:
        """Confirm write operations and release pages."""
        with self.lock:
            for key, page_index in written_keys_to_confirm:
                self.key_to_index[key] = page_index

            for page_index in pages_to_release:
                if page_index not in self.free_pages:
                    self.free_pages.append(page_index)

    def delete_keys(self, keys: List[str]) -> int:
        """Delete keys and return count of deleted keys."""
        with self.lock:
            count = 0
            for key in keys:
                if key in self.key_to_index:
                    page_index = self.key_to_index.pop(key)
                    if page_index not in self.free_pages:
                        self.free_pages.append(page_index)
                    count += 1
            return count

    def clear_all(self) -> None:
        """Clear all metadata."""
        with self.lock:
            self.free_pages = list(range(self.num_pages))
            self.key_to_index.clear()

    def get_page_indices(self, keys: List[str]) -> List[Optional[int]]:
        """Get page indices for keys."""
        with self.lock:
            return [self.key_to_index.get(key) for key in keys]


class GlobalMetadataState:
    """Manages the state for all ranks and persistence."""

    def __init__(self, persistence_path: Optional[str], save_interval: int):
        self.global_lock = threading.RLock()
        self.ranks: Dict[int, RankMetadata] = {}
        self.persistence_path = Path(persistence_path) if persistence_path else None
        self.save_interval = save_interval
        self.save_timer: Optional[threading.Timer] = None
        self.is_shutting_down = False

    def load_from_disk(self):
        if not self.persistence_path or not self.persistence_path.exists():
            logging.info("Persistence file not found. Starting with a clean state.")
            return

        logging.info(f"Loading state from {self.persistence_path}")
        try:
            with open(self.persistence_path, "r") as f:
                persisted_data = json.load(f)

            with self.global_lock:
                for rank_id_str, data in persisted_data.items():
                    rank_id = int(rank_id_str)
                    num_pages = data["num_pages"]
                    rank_meta = RankMetadata(num_pages)
                    rank_meta.free_pages = data["free_pages"]
                    rank_meta.key_to_index = dict(data["key_to_index"])
                    self.ranks[rank_id] = rank_meta
                logging.info(
                    f"Successfully loaded metadata for {len(self.ranks)} ranks."
                )
        except (json.JSONDecodeError, KeyError, TypeError) as e:
            logging.error(
                f"Failed to load or parse persistence file: {e}. Starting fresh.",
                exc_info=True,
            )
            self.ranks.clear()

    def save_to_disk(self):
        if not self.persistence_path:
            return

        logging.info("Persisting metadata to disk...")
        with self.global_lock:
            serializable_state = {}
            for rank_id, rank_meta in self.ranks.items():
                with rank_meta.lock:
                    serializable_state[rank_id] = {
                        "num_pages": rank_meta.num_pages,
                        "free_pages": rank_meta.free_pages,
                        "key_to_index": list(rank_meta.key_to_index.items()),
                    }

        try:
            temp_path = self.persistence_path.with_suffix(".tmp")
            with open(temp_path, "w") as f:
                json.dump(serializable_state, f, indent=4)
            temp_path.rename(self.persistence_path)
            logging.info(f"Metadata successfully persisted to {self.persistence_path}")
        except Exception as e:
            logging.error(f"Failed to save metadata to disk: {e}", exc_info=True)

    def schedule_save(self):
        if self.is_shutting_down or not self.persistence_path:
            return
        self.save_to_disk()
        self.save_timer = threading.Timer(self.save_interval, self.schedule_save)
        self.save_timer.start()

    def shutdown(self):
        logging.info("Shutting down metadata server...")
        self.is_shutting_down = True
        if self.save_timer:
            self.save_timer.cancel()
        self.save_to_disk()
        logging.info("Shutdown complete.")


# --- Global MetadataServer implementation ---
class Hf3fsMetadataServer:
    """HF3FS Metadata Server that manages metadata for multiple ranks."""

    def __init__(self, persistence_path: Optional[str] = None, save_interval: int = 60):
        self.state = GlobalMetadataState(persistence_path, save_interval)
        self.app = FastAPI()
        self._setup_routes()

    def _setup_routes(self):
        """Setup FastAPI routes."""
        self.app.post("/{rank}/initialize")(self.initialize)
        self.app.post("/{rank}/exists")(self.exists)
        self.app.post("/{rank}/reserve_and_allocate_page_indices")(
            self.reserve_and_allocate_page_indices
        )
        self.app.post("/{rank}/confirm_write")(self.confirm_write)
        self.app.post("/{rank}/delete_keys")(self.delete_keys)
        self.app.post("/{rank}/clear")(self.clear)
        self.app.post("/{rank}/get_page_indices")(self.get_page_indices)

    def get_rank_metadata(self, rank: int) -> RankMetadata:
        """Get rank metadata with proper error handling."""
        with self.state.global_lock:
            if rank not in self.state.ranks:
                raise HTTPException(
                    status_code=404,
                    detail=f"Rank {rank} not initialized. Please call /{{rank}}/initialize first.",
                )
            return self.state.ranks[rank]

    async def initialize(self, rank: int, request: Request):
        """Initialize a rank with specified number of pages."""
        data = await request.json()
        num_pages = data["num_pages"]
        with self.state.global_lock:
            if rank in self.state.ranks:
                logging.info(
                    f"Rank {rank} already exists. Initialization request ignored."
                )
                if self.state.ranks[rank].num_pages != num_pages:
                    logging.warning(
                        f"Rank {rank} initialized with different num_pages. Existing: {self.state.ranks[rank].num_pages}, New: {num_pages}"
                    )
            else:
                logging.info(f"Initializing new Rank {rank} with {num_pages} pages.")
                self.state.ranks[rank] = RankMetadata(num_pages)
        return {"message": f"Rank {rank} is ready."}

    async def exists(self, rank: int, request: Request):
        """Check if keys exist in metadata."""
        data = await request.json()
        keys = data["keys"]
        metadata = self.get_rank_metadata(rank)
        results = metadata.exists_keys(keys)
        return {"exists": results}

    async def reserve_and_allocate_page_indices(self, rank: int, request: Request):
        """Reserve and allocate page indices for keys."""
        data = await request.json()
        metadata = self.get_rank_metadata(rank)
        keys = data["keys"]
        results = metadata.reserve_and_allocate_page_indices(keys)
        return {"indices": results}

    async def confirm_write(self, rank: int, request: Request):
        """Confirm write operations and release pages."""
        data = await request.json()
        metadata = self.get_rank_metadata(rank)
        success_written_keys = data.get("written_keys_to_confirm", [])
        released_pages = data.get("pages_to_release", [])

        metadata.confirm_write(success_written_keys, released_pages)

        return {
            "message": f"Rank {rank}: Write confirmed for {len(success_written_keys)} keys. {len(released_pages)} pages released."
        }

    async def delete_keys(self, rank: int, request: Request):
        """Delete keys from metadata."""
        data = await request.json()
        metadata = self.get_rank_metadata(rank)
        count = metadata.delete_keys(data["keys"])
        return {"message": f"Rank {rank}: {count} keys deleted."}

    async def clear(self, rank: int):
        """Clear all metadata for a rank."""
        metadata = self.get_rank_metadata(rank)
        metadata.clear_all()
        return {"message": f"Rank {rank}: Metadata cleared."}

    async def get_page_indices(self, rank: int, request: Request):
        """Get page indices for keys."""
        data = await request.json()
        metadata = self.get_rank_metadata(rank)
        keys = data["keys"]
        results = metadata.get_page_indices(keys)
        return {"indices": results}

    def run(self, host: str = "0.0.0.0", port: int = 18000):
        """Run the metadata server."""
        self.state.load_from_disk()
        if self.state.persistence_path:
            self.state.schedule_save()
            atexit.register(self.state.shutdown)

        import uvicorn

        logging.info(f"Starting metadata server on http://{host}:{port}")
        if self.state.persistence_path:
            logging.info(
                f"Persistence is ENABLED. Saving to '{self.state.persistence_path}' every {self.state.save_interval} seconds."
            )
        else:
            logging.info("Persistence is DISABLED.")

        uvicorn.run(self.app, host=host, port=port)


# --- Client implementation ---
class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
    """Global http metadata client for HF3FS."""

    def __init__(self, base_url: str, max_retries: int = 3):
        self.base_url = base_url.rstrip("/")
        self._session = requests.Session()

        retry_strategy = Retry(
            total=max_retries,
            backoff_factor=0.3,
            status_forcelist=[500, 502, 503, 504],
            allowed_methods=["GET", "POST"],
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self._session.mount("http://", adapter)

    def _post(self, endpoint: str, json_data: dict) -> dict:
        try:
            response = self._session.post(f"{self.base_url}/{endpoint}", json=json_data)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            logging.error(f"Failed to POST to {endpoint} after retries: {e}")
            raise RuntimeError(f"Failed to connect to metadata server: {e}") from e

    def initialize(self, rank: int, num_pages: int) -> None:
        self._post(f"{rank}/initialize", {"num_pages": num_pages})

    def reserve_and_allocate_page_indices(
        self, rank: int, keys: List[Tuple[str, str]]
    ) -> List[Tuple[bool, int]]:
        response = self._post(
            f"{rank}/reserve_and_allocate_page_indices", {"keys": keys}
        )
        return [tuple(item) for item in response.get("indices")]

    def confirm_write(
        self,
        rank: int,
        written_keys_to_confirm: List[Tuple[str, int]],
        pages_to_release: List[int],
    ) -> None:
        self._post(
            f"{rank}/confirm_write",
            {
                "written_keys_to_confirm": written_keys_to_confirm,
                "pages_to_release": pages_to_release,
            },
        )

    def delete_keys(self, rank: int, keys: List[str]) -> None:
        self._post(f"{rank}/delete_keys", {"keys": keys})

    def exists(self, rank: int, keys: List[str]) -> List[bool]:
        response = self._post(f"{rank}/exists", {"keys": keys})
        return response.get("exists", [])

    def clear(self, rank: int) -> None:
        self._post(f"{rank}/clear", {})

    def get_page_indices(self, rank: int, keys: List[str]) -> List[Optional[int]]:
        response = self._post(f"{rank}/get_page_indices", {"keys": keys})
        return response.get("indices")


class Hf3fsLocalMetadataClient(Hf3fsMetadataInterface):
    """Local metadata client that directly operates on single RankMetadata in memory without metadata server."""

    def __init__(self):
        self.rank_metadata = None

    def initialize(self, rank: int, num_pages: int) -> None:
        self.rank_metadata = RankMetadata(num_pages)

    def reserve_and_allocate_page_indices(
        self, rank: int, keys: List[Tuple[str, str]]
    ) -> List[Tuple[bool, int]]:
        """Reserve and allocate page indices for keys."""
        return self.rank_metadata.reserve_and_allocate_page_indices(keys)

    def confirm_write(
        self,
        rank: int,
        written_keys_to_confirm: List[Tuple[str, int]],
        pages_to_release: List[int],
    ) -> None:
        """Confirm write operations."""
        self.rank_metadata.confirm_write(written_keys_to_confirm, pages_to_release)

    def delete_keys(self, rank: int, keys: List[str]) -> None:
        """Delete keys."""
        self.rank_metadata.delete_keys(keys)

    def exists(self, rank: int, keys: List[str]) -> List[bool]:
        """Check if keys exist."""
        return self.rank_metadata.exists_keys(keys)

    def clear(self, rank: int) -> None:
        """Clear all metadata for rank."""
        self.rank_metadata.clear_all()

    def get_page_indices(self, rank: int, keys: List[str]) -> List[Optional[int]]:
        """Get page indices for keys."""
        return self.rank_metadata.get_page_indices(keys)


def run_metadata_server(
    host: str = "0.0.0.0",
    port: int = 18000,
    persistence_path: Optional[str] = None,
    save_interval: int = 60,
):
    """Run the HF3FS metadata server."""
    global server
    server = Hf3fsMetadataServer(
        persistence_path=persistence_path, save_interval=save_interval
    )

    server.run(host=host, port=port)


# --- Main Execution ---
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="HF3FS Metadata Server")
    parser.add_argument(
        "--host", type=str, default="0.0.0.0", help="Host to bind the server to."
    )
    parser.add_argument(
        "--port", type=int, default=18000, help="Port to run the server on."
    )
    parser.add_argument(
        "--persistence-path",
        type=str,
        default=None,
        help="Path to the file for persisting metadata. If not provided, persistence is disabled.",
    )
    parser.add_argument(
        "--save-interval",
        type=int,
        default=60,
        help="Interval in seconds for periodically saving metadata to disk.",
    )
    args = parser.parse_args()

    run_metadata_server(args.host, args.port, args.persistence_path, args.save_interval)
