import os
from dataclasses import dataclass
from typing import List


@dataclass
class NPGraphMetadata:
    """A Graph Metadata returned by the Neuronpedia API."""

    id: str  # A unique ID automatically generated by Neuronpedia
    model_id: str  # The model that the graph belongs to
    slug: str  # A user-defined identifier for the graph. Must be globally unique combined with model_id.
    prompt_tokens: List[str]  # The tokens that were used to generate the graph
    prompt: str  # The prompt that was used to generate the graph
    title_prefix: str
    json_url: str  # URL to the actual graph JSON file
    url: str | None = None  # URL to the graph on Neuronpedia
    url_embed: str | None = None  # URL to the embedded graph on Neuronpedia

    def __post_init__(self):
        if not self.url:
            USE_LOCALHOST = os.getenv("USE_LOCALHOST", "false").lower() == "true"
            BASE_URL = (
                "https://neuronpedia.org"
                if not USE_LOCALHOST
                else "http://localhost:3000"
            )
            self.url = f"{BASE_URL}/{self.model_id}/graph?slug={self.slug}"
            self.url_embed = f"{self.url}&embed=true"

    def __eq__(self, other: "NPGraphMetadata") -> bool:
        return (
            self.model_id == other.model_id
            and self.slug == other.slug
            and self.prompt_tokens == other.prompt_tokens
            and self.prompt == other.prompt
            and self.title_prefix == other.title_prefix
            and self.json_url == other.json_url
            and self.url == other.url
        )

    def delete(self):
        # import here to avoid circular import
        from neuronpedia.requests.graph_request import GraphRequest

        return GraphRequest().delete(self)

    @classmethod
    def get(cls, model_id: str, slug: str) -> "NPGraphMetadata":
        # import here to avoid circular import
        from neuronpedia.requests.graph_request import GraphRequest

        return GraphRequest().get(model_id, slug)

    @classmethod
    def list_owned(cls) -> List["NPGraphMetadata"]:
        # import here to avoid circular import
        from neuronpedia.requests.graph_request import GraphRequest

        return GraphRequest().list_owned()

    @classmethod
    def generate(
        cls,
        model_id: str,
        prompt: str,
        graph_id: str,
        max_n_logits: int = 10,
        desired_logit_prob: float = 0.95,
        node_threshold: float = 0.8,
        edge_threshold: float = 0.98,
    ) -> "NPGraphMetadata":
        # import here to avoid circular import
        from neuronpedia.requests.graph_request import GraphRequest

        return GraphRequest().generate(
            model_id,
            prompt,
            graph_id,
            max_n_logits,
            desired_logit_prob,
            node_threshold,
            edge_threshold,
        )

    @classmethod
    def upload(cls, json_str: str) -> "NPGraphMetadata":
        # import here to avoid circular import
        from neuronpedia.requests.graph_request import GraphRequest

        return GraphRequest().upload(json_str)

    @classmethod
    def upload_file(cls, filepath: str) -> "NPGraphMetadata":
        # import here to avoid circular import
        from neuronpedia.requests.graph_request import GraphRequest

        return GraphRequest().upload_file(filepath)
