import json
import logging
from pathlib import Path
from typing import List, Dict, Optional, Union, Tuple
from collections import deque

import torch
import numpy as np
from e5_encoder import E5Encoder

DEFAULT_PLAN = ["explain", "support", "respond", "help"]
SPAM_REPEATS = 3
SIM_THRESHOLD = 0.7

class IntentProcessor:
    VOCAB_FILE = "behavioral_vocab.json"

    def __init__(
        self,
        role_text: str,
        top_k: int = 7,
        vocab_path: Optional[str] = None,
        encoder: Optional[E5Encoder] = None,
        verbose: bool = False
    ):
        self.logger = logging.getLogger(__name__)
        if not self.logger.handlers:
            logging.basicConfig(level=logging.INFO)

        self.role_text = role_text.strip()[:1000]
        self.top_k = min(max(top_k, 1), 10)
        self.verbose = verbose

        self.encoder = encoder or self._init_encoder()
        self.vocab: List[str] = []
        self.vocab_embeddings: Optional[torch.Tensor] = None
        self.history = deque(maxlen=100)

        self._last_plan: List[str] = []
        self._last_spam_flag: bool = False

        self._load_vocab(vocab_path)

    def _log(self, msg: str):
        if self.verbose:
            self.logger.info(f"[IntentProcessor] {msg}")

    def _init_encoder(self) -> Optional[E5Encoder]:
        try:
            enc = E5Encoder()
            self._log(f"E5Encoder загружен на {enc.device}")
            return enc
        except Exception as e:
            self._log(f"Ошибка загрузки E5Encoder: {e}")
            return None

    def _load_vocab(self, vocab_path: Optional[str]):
        try:
            path = Path(vocab_path) if vocab_path else Path(__file__).resolve().parent / self.VOCAB_FILE
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            self.vocab = data.get("verbs", DEFAULT_PLAN)
            if self.encoder:
                self.vocab_embeddings = torch.stack([
                    torch.tensor(self.encoder.encode(word)) for word in self.vocab
                ])
                self._log(f"Вокабуляр загружен: {len(self.vocab)} терминов")
        except Exception as e:
            self._log(f"Ошибка загрузки вокабуляра: {e}")
            self.vocab = DEFAULT_PLAN

    def generate_plan(self, return_metadata: bool = False) -> Union[List[str], Tuple[List[str], Dict]]:
        self._last_plan = []
        self._last_spam_flag = False

        if not self.role_text or not self.encoder or self.vocab_embeddings is None:
            self._log("⚠️ Используется дефолтный план")
            self._last_plan = DEFAULT_PLAN.copy()
            return self._return(self._last_plan, metadata=return_metadata)

        try:
            role_vec_np = self.encoder.encode(self.role_text)
            role_vec = torch.tensor(role_vec_np).unsqueeze(0)

            if self._is_spam(role_vec, self.role_text):
                self._last_spam_flag = True
                self._last_plan = ["reject_spam"]
                return self._return(self._last_plan, metadata=return_metadata)

            scores = torch.nn.functional.cosine_similarity(role_vec, self.vocab_embeddings, dim=1)
            max_score = float(scores.max().item())

            if max_score < SIM_THRESHOLD:
                self._last_plan = ["reject_prompt"]
            else:
                top_indices = torch.topk(scores, k=self.top_k).indices.tolist()
                self._last_plan = [self.vocab[i] for i in top_indices]

            self._log(f"📜 План: {self._last_plan}")
            self.history.append({"text": self.role_text, "vec": role_vec.squeeze(0).numpy()})

            return self._return(self._last_plan, {
                "spam": self._last_spam_flag,
                "score_max": max_score
            } if return_metadata else None)

        except Exception as e:
            self._log(f"Ошибка генерации плана: {e}")
            self._last_plan = DEFAULT_PLAN.copy()
            return self._return(self._last_plan, metadata=return_metadata)

    def _return(self, plan: List[str], metadata: Optional[Dict] = None):
        return (plan, metadata) if metadata is not None else plan

    def _is_spam(self, vec: torch.Tensor, text: str) -> bool:
        repeat_count = sum(1 for p in self.history if p["text"] == text)
        if repeat_count > SPAM_REPEATS:
            self._log("⚠️ Повторяющийся текст — spam")
            return True

        for prev in self.history:
            prev_vec = torch.tensor(prev["vec"])
            if torch.nn.functional.cosine_similarity(vec, prev_vec.unsqueeze(0)).item() > 0.91:
                self._log("⚠️ Семантический повтор — spam")
                return True
        return False

    def diagnostics(self) -> Dict:
        return {
            "available": bool(self.encoder and self.vocab_embeddings is not None),
            "vocab_size": len(self.vocab),
            "top_k": self.top_k,
            "role_text": self.role_text,
            "last_plan": self._last_plan,
            "spam_flag": self._last_spam_flag,
            "history_size": len(self.history),
            "device": self.encoder.device if self.encoder else "none"
        }