"""Knowledge Graph Tool for querying biomedical knowledge graphs."""

from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import aiohttp


@dataclass
class Entity:
    """Base entity from knowledge graph."""

    id: str
    name: str
    entity_type: str

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Entity":
        return cls(
            id=str(data.get("id", "")),
            name=data.get("name") or data.get("label") or data.get("symbol", ""),
            entity_type=data.get("type", "unknown"),
        )


@dataclass
class Edge:
    """Relationship edge in knowledge graph."""

    source_id: str
    source_type: str
    target_id: str
    target_type: str
    relation_type: str

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Edge":
        return cls(
            source_id=str(data.get("source_id", "")),
            source_type=data.get("source_type", ""),
            target_id=str(data.get("target_id", "")),
            target_type=data.get("target_type", ""),
            relation_type=data.get("relation_type", ""),
        )


@dataclass
class SimilarDisease:
    """Disease similarity result."""

    disease_id: str
    disease_name: str
    similarity_score: float
    shared_count: int

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SimilarDisease":
        return cls(
            disease_id=str(data.get("disease_id", "")),
            disease_name=data.get("disease_name", ""),
            similarity_score=float(data.get("similarity_score", 0)),
            shared_count=int(data.get("shared_count", 0)),
        )


@dataclass
class PathStep:
    """Step in shortest path."""

    node_id: str
    node_name: str
    node_type: str
    relation: Optional[str] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "PathStep":
        return cls(
            node_id=str(data.get("node_id", "")),
            node_name=data.get("node_name", ""),
            node_type=data.get("node_type", ""),
            relation=data.get("relation"),
        )


class KnowledgeGraphTool:
    """
    Tool for querying biomedical knowledge graphs.

    Usage:
        # Initialize with gateway URL
        kg = KnowledgeGraphTool(base_url="http://gateway:8080/s/knowledge-graph")

        # Or from gateway client
        kg = KnowledgeGraphTool.from_gateway(gateway_client)

        # Query diseases
        diseases = await kg.search_diseases(name="Huntington")

        # Find similar diseases
        similar = await kg.find_similar_diseases("Huntington Disease", limit=10)

        # Get entity connections
        edges = await kg.get_entity_network("BRCA1", "protein")

        # Find path between entities
        path = await kg.find_shortest_path("BRCA1", "protein", "Breast Cancer", "disease")
    """

    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self._session: Optional[aiohttp.ClientSession] = None

    @classmethod
    def from_gateway(cls, gateway_client) -> "KnowledgeGraphTool":
        """Create from GatewayClient instance."""
        return cls(gateway_client.knowledge_graph_url)

    async def _get_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            timeout = aiohttp.ClientTimeout(total=30.0)
            self._session = aiohttp.ClientSession(timeout=timeout)
        return self._session

    async def close(self):
        """Close the HTTP session."""
        if self._session and not self._session.closed:
            await self._session.close()

    async def _get(self, path: str, params: Optional[Dict] = None) -> Any:
        session = await self._get_session()
        url = f"{self.base_url}{path}"
        async with session.get(url, params=params) as response:
            response.raise_for_status()
            return await response.json()

    # === Entity Lookup Methods ===

    async def search_diseases(
        self,
        name: Optional[str] = None,
        orphacode: Optional[str] = None,
        omim: Optional[str] = None,
        limit: int = 20,
    ) -> List[Dict]:
        """Search for diseases by name or codes."""
        params: Dict[str, Any] = {"limit": limit}
        if name:
            params["label"] = name
        if orphacode:
            params["orphacode"] = orphacode
        if omim:
            params["omim"] = omim
        return await self._get("/diseases/", params)

    async def get_disease(self, disease_id: int) -> Dict:
        """Get disease by ID."""
        return await self._get(f"/diseases/{disease_id}")

    async def search_proteins(
        self,
        symbol: Optional[str] = None,
        ensembl_id: Optional[str] = None,
        limit: int = 20,
    ) -> List[Dict]:
        """Search for proteins/genes."""
        params: Dict[str, Any] = {"limit": limit}
        if symbol:
            params["symbol"] = symbol
        if ensembl_id:
            params["ensembl_id"] = ensembl_id
        return await self._get("/proteins/", params)

    async def search_phenotypes(
        self,
        name: Optional[str] = None,
        hpo_id: Optional[str] = None,
        limit: int = 20,
    ) -> List[Dict]:
        """Search for phenotypes (HPO terms)."""
        params: Dict[str, Any] = {"limit": limit}
        if name:
            params["name"] = name
        if hpo_id:
            params["hpo_id"] = hpo_id
        return await self._get("/phenotypes/", params)

    async def search_drugs(
        self,
        name: Optional[str] = None,
        limit: int = 20,
    ) -> List[Dict]:
        """Search for drugs."""
        params: Dict[str, Any] = {"limit": limit}
        if name:
            params["name"] = name
        return await self._get("/drugs/", params)

    async def search_pathways(
        self,
        name: Optional[str] = None,
        limit: int = 20,
    ) -> List[Dict]:
        """Search for biological pathways."""
        params: Dict[str, Any] = {"limit": limit}
        if name:
            params["name"] = name
        return await self._get("/pathways/", params)

    # === Similarity Analysis ===

    async def find_similar_diseases(
        self,
        disease_name: str,
        limit: int = 10,
    ) -> List[SimilarDisease]:
        """Find diseases similar by shared phenotypes."""
        data = await self._get(
            "/diseases/similarity",
            params={"disease_name": disease_name, "limit": limit},
        )
        return [SimilarDisease.from_dict(d) for d in data.get("similar_diseases", [])]

    async def find_similar_diseases_by_genes(
        self,
        disease_name: str,
        limit: int = 10,
    ) -> List[SimilarDisease]:
        """Find diseases similar by shared genes."""
        data = await self._get(
            "/diseases/similarity-by-genes",
            params={"disease_name": disease_name, "limit": limit},
        )
        return [SimilarDisease.from_dict(d) for d in data.get("similar_diseases", [])]

    # === Graph Traversal ===

    async def get_entity_network(
        self,
        entity_name: str,
        entity_type: str,
    ) -> List[Edge]:
        """
        Get all connections for an entity.

        Args:
            entity_name: Name of the entity (e.g., "BRCA1", "Huntington Disease")
            entity_type: One of: protein, phenotype, disease, pathway,
                        molecular_function, cellular_component, biological_process
        """
        data = await self._get(
            "/edges/network",
            params={"entity_name": entity_name, "entity_type": entity_type},
        )
        return [Edge.from_dict(e) for e in data]

    async def k_hop_traversal(
        self,
        entity_name: str,
        entity_type: str,
        k: int = 2,
        limit_edges: int = 100,
    ) -> Dict:
        """
        Perform k-hop BFS traversal from entity.

        Returns subgraph with nodes and edges within k hops.
        """
        return await self._get(
            "/graph/k-hop",
            params={
                "entity_name": entity_name,
                "entity_type": entity_type,
                "k": k,
                "limit_edges": limit_edges,
            },
        )

    async def find_shortest_path(
        self,
        start_name: str,
        start_type: str,
        end_name: str,
        end_type: str,
        max_hops: int = 5,
    ) -> List[PathStep]:
        """
        Find shortest path between two entities (Mechanism of Action).

        Args:
            start_name: Starting entity name
            start_type: Starting entity type
            end_name: Target entity name
            end_type: Target entity type
            max_hops: Maximum path length
        """
        data = await self._get(
            "/paths/shortest",
            params={
                "start_name": start_name,
                "start_type": start_type,
                "end_name": end_name,
                "end_type": end_type,
                "max_hops": max_hops,
            },
        )
        return [PathStep.from_dict(s) for s in data.get("path", [])]
