"""SDK Service Extensions — Graph, Vector, Training, Streams"""

import logging
from dataclasses import dataclass
from typing import List, Optional

logger = logging.getLogger(__name__)


@dataclass
class SimilarTrajectoryResult:
    trajectory_id: str
    score: float
    task_type: str
    outcome: str
    skills_used: list
    metadata: dict


@dataclass
class GraphRecommendation:
    pattern_id: str
    pattern_name: str
    success_rate: float
    skills_used: list
    confidence: float


@dataclass
class SkillGap:
    skill_id: str
    skill_name: str
    domain: str
    description: str


@dataclass
class TrainingStats:
    is_available: bool
    total_trajectories_buffered: int
    training_runs: int
    last_training_metrics: Optional[dict]


@dataclass
class ServiceHealth:
    neo4j: bool
    pinecone: bool
    redis_streams: bool
    grpo: bool


class ServiceMixin:
    """
    Service methods for ZubblClient.

    Service methods for ZubblClient (graph, vector, training, streams).

    These methods use self._sync_client (httpx.Client) and self._client
    (httpx.AsyncClient) already available on ZubblClient.

    All methods degrade gracefully if services are unavailable.
    """

    # --- Graph-Based Recommendations ---

    def get_graph_recommendations(
        self, task_type: str, min_confidence: float = 0.5, limit: int = 5
    ) -> List[GraphRecommendation]:
        """
        Get graph-based pattern recommendations for a task type.
        Uses Neo4j skill graph to find successful patterns.
        """
        try:
            resp = self._sync_client.get(
                f"/graph/recommendations/{task_type}",
                params={"min_confidence": min_confidence, "limit": limit},
            )
            resp.raise_for_status()
            data = resp.json()
            return [GraphRecommendation(**r) for r in data.get("recommendations", [])]
        except Exception as e:
            logger.warning(f"Graph recommendations unavailable: {e}")
            return []

    async def get_graph_recommendations_async(
        self, task_type: str, min_confidence: float = 0.5, limit: int = 5
    ) -> List[GraphRecommendation]:
        """Async version of get_graph_recommendations."""
        try:
            resp = await self._client.get(
                f"/graph/recommendations/{task_type}",
                params={"min_confidence": min_confidence, "limit": limit},
            )
            resp.raise_for_status()
            data = resp.json()
            return [GraphRecommendation(**r) for r in data.get("recommendations", [])]
        except Exception as e:
            logger.warning(f"Graph recommendations unavailable: {e}")
            return []

    def get_skill_gaps(self, agent_id: str, task_type: str) -> List[SkillGap]:
        """Find skills the agent is missing for a task type."""
        try:
            resp = self._sync_client.get(
                f"/graph/skill-gaps/{agent_id}/{task_type}",
            )
            resp.raise_for_status()
            data = resp.json()
            return [SkillGap(**s) for s in data.get("skill_gaps", [])]
        except Exception as e:
            logger.warning(f"Skill gap analysis unavailable: {e}")
            return []

    # --- Vector Similarity Search ---

    def find_similar_trajectories(
        self, query_text: str, top_k: int = 5, status_filter: str = None
    ) -> List[SimilarTrajectoryResult]:
        """
        Find similar past trajectories using vector similarity.
        Uses Pinecone for fast semantic search.
        """
        params = {"text": query_text, "top_k": top_k}
        if status_filter:
            params["status_filter"] = status_filter
        try:
            resp = self._sync_client.get(
                "/vector/similar",
                params=params,
            )
            resp.raise_for_status()
            data = resp.json()
            return [SimilarTrajectoryResult(**t) for t in data.get("results", [])]
        except Exception as e:
            logger.warning(f"Vector search unavailable: {e}")
            return []

    async def find_similar_trajectories_async(
        self, query_text: str, top_k: int = 5, status_filter: str = None
    ) -> List[SimilarTrajectoryResult]:
        """Async version of find_similar_trajectories."""
        params = {"text": query_text, "top_k": top_k}
        if status_filter:
            params["status_filter"] = status_filter
        try:
            resp = await self._client.get(
                "/vector/similar",
                params=params,
            )
            resp.raise_for_status()
            data = resp.json()
            return [SimilarTrajectoryResult(**t) for t in data.get("results", [])]
        except Exception as e:
            logger.warning(f"Vector search unavailable: {e}")
            return []

    def find_successful_patterns(
        self, task_description: str, top_k: int = 5, min_success_rate: float = 0.7
    ) -> List[dict]:
        """Find similar successful patterns via vector search."""
        try:
            resp = self._sync_client.get(
                "/vector/patterns",
                params={
                    "text": task_description,
                    "top_k": top_k,
                    "min_success_rate": min_success_rate,
                },
            )
            resp.raise_for_status()
            return resp.json().get("patterns", [])
        except Exception as e:
            logger.warning(f"Pattern search unavailable: {e}")
            return []

    # --- Training ---

    def get_training_stats(self) -> Optional[TrainingStats]:
        """Get GRPO training statistics."""
        try:
            resp = self._sync_client.get(
                "/training/stats",
            )
            resp.raise_for_status()
            return TrainingStats(**resp.json())
        except Exception as e:
            logger.warning(f"Training stats unavailable: {e}")
            return None

    def trigger_training(self) -> Optional[dict]:
        """Manually trigger GRPO training."""
        try:
            resp = self._sync_client.post(
                "/training/trigger",
                timeout=30.0,
            )
            resp.raise_for_status()
            return resp.json()
        except Exception as e:
            logger.warning(f"Training trigger unavailable: {e}")
            return None

    # --- Services Health ---

    def get_services_health(self) -> ServiceHealth:
        """Check which services are available."""
        try:
            resp = self._sync_client.get(
                "/services/health",
            )
            resp.raise_for_status()
            data = resp.json()
            return ServiceHealth(**data)
        except Exception as e:
            logger.warning(f"Services health check failed: {e}")
            return ServiceHealth(neo4j=False, pinecone=False, redis_streams=False, grpo=False)

    # --- Streams Info ---

    def get_streams_info(self) -> Optional[dict]:
        """Get Redis Streams queue information."""
        try:
            resp = self._sync_client.get(
                "/streams/info",
            )
            resp.raise_for_status()
            return resp.json()
        except Exception as e:
            logger.warning(f"Streams info unavailable: {e}")
            return None
