from __future__ import annotations

import logging
from collections import defaultdict
from typing import Any, Optional

import numpy as np

from app.services.embeddings_service import EmbeddingService

_logger = logging.getLogger(__name__)

_DIMENSION = 384


class SemanticRouterService:
    def __init__(self, embedding_service: EmbeddingService) -> None:
        self._embedding_service = embedding_service

    def route(
        self,
        query: str,
        routes: list[dict[str, Any]],
        threshold: float = 0.7,
    ) -> dict[str, Any]:
        if not query or not query.strip():
            return {"success": False, "name": None, "models": [], "error": "Query is empty."}
        if not routes:
            return {"success": False, "name": None, "models": [], "error": "No routes provided."}

        if not self._embedding_service.is_loaded(_DIMENSION):
            return {"success": False, "name": None, "models": [], "error": "Embedding model not loaded."}

        all_utterances: list[str] = []
        utterance_to_route: list[int] = []

        for idx, route in enumerate(routes):
            utterances = route.get("utterances", [])
            if not utterances:
                continue
            all_utterances.extend(utterances)
            utterance_to_route.extend([idx] * len(utterances))

        if not all_utterances:
            return {"success": False, "name": None, "models": [], "error": "No utterances found in any route."}

        try:
            all_texts = all_utterances + [query]
            embeddings = self._embedding_service.generate_embedding(all_texts, _DIMENSION)
        except Exception as exc:
            _logger.error("Embedding failed: %s", exc)
            return {"success": False, "name": None, "models": [], "error": str(exc)}

        utterance_embs = np.array(embeddings[:-1], dtype=np.float32)
        query_emb = np.array(embeddings[-1], dtype=np.float32).reshape(1, -1)

        similarities = (utterance_embs @ query_emb.T).flatten()
        best_idx = int(np.argmax(similarities))
        best_score = float(similarities[best_idx])

        # Route-level scoring: take the max score per route
        route_scores: dict[int, list[float]] = defaultdict(list)
        for i, score in enumerate(similarities):
            route_scores[utterance_to_route[i]].append(float(score))

        route_best_scores = {rid: max(scores) for rid, scores in route_scores.items()}
        sorted_routes = sorted(route_best_scores.items(), key=lambda x: x[1], reverse=True)
        best_route_score = sorted_routes[0][1]
        second_best_route_score = sorted_routes[1][1] if len(sorted_routes) > 1 else 1.0
        route_margin = best_route_score - second_best_route_score

        if best_score < threshold or route_margin < 0.001:
            return {
                "success": True,
                "name": None,
                "models": [],
                "error": None,
                "confidence": best_score,
                "margin": route_margin,
                "threshold": threshold,
                "matched_utterance": all_utterances[best_idx],
            }

        matched_route_idx = utterance_to_route[best_idx]
        matched_route = routes[matched_route_idx]

        return {
            "success": True,
            "name": matched_route.get("name"),
            "models": matched_route.get("models", []),
            "error": None,
            "confidence": best_score,
            "margin": route_margin,
            "threshold": threshold,
            "matched_utterance": all_utterances[best_idx],
        }