import asyncio
import logging
import torch
from difflib import SequenceMatcher
from time import time
from typing import Dict, List, Tuple

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

from config_settings_public import MODEL_PATHS, MODEL_CONFIG, MODEL_WEIGHTS
from model_manager_public import ModelManager
from response_optimizer_public import ResponseOptimizer

logger = logging.getLogger(__name__)


class QueryProcessor:
    def __init__(self, model_manager: ModelManager, response_optimizer: ResponseOptimizer):
        self.model_manager = model_manager
        self.response_optimizer = response_optimizer
        self.model_paths = MODEL_PATHS
        self.model_config = MODEL_CONFIG
        self.memory_bank = []

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            torch.cuda.set_device(0)
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.enabled = True
            torch.set_float32_matmul_precision('high')
            logger.info(f"✅ QueryProcessor initialized on {self.device}")
        else:
            logger.warning("⚠️ Running on CPU")

        self.model_tiers = {
            "balanced": ["gemma", "mistral", "llama"]
        }

    def select_models(self, query: str) -> List[str]:
        return [m for m in self.model_tiers["balanced"] if m in self.model_paths]

    async def process_query_with_fusion(self, query: str) -> str:
        selected_models = self.select_models(query)
        logger.info(f"🧠 Selected models: {selected_models}")

        tasks = [asyncio.create_task(self.invoke_model(m, query)) for m in selected_models]
        responses = await asyncio.gather(*tasks, return_exceptions=True)

        valid_responses = {
            m: r for m, r in zip(selected_models, responses)
            if isinstance(r, str) and r.strip()
        }

        if not valid_responses:
            logger.error("❌ All models failed or returned empty responses.")
            return "⚠️ No valid responses received from models."

        logger.info(f"✅ {len(valid_responses)} valid responses received.")
        fused = await self.advanced_fusion(valid_responses, query)
        self.memory_bank.append({"query": query, "response": fused, "timestamp": time()})
        return fused

    async def invoke_model(self, model_name: str, query: str) -> str:
        try:
            model = self.model_manager.loaded_models.get(model_name)
            if not model:
                model = self.model_manager.load_model_from_disk(model_name)
                if not model:
                    raise RuntimeError(f"Model '{model_name}' failed to load.")

            config = self.model_config.get(model_name, {})
            max_tokens = config.get("max_tokens", 800)
            temperature = config.get("temperature", 0.7)

            logger.info(f"🔧 Invoking '{model_name}' with temp={temperature}, max_tokens={max_tokens}")

            with torch.no_grad():
                completion = await asyncio.to_thread(
                    model.create_completion,
                    prompt=query,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    stop=["</s>", "User:", "Assistant:"]
                )

            response_text = completion.get("choices", [{}])[0].get("text", "").strip()
            return response_text

        except Exception as e:
            logger.error(f"❌ Error in model '{model_name}': {e}")
            return ""

    def clean_response(self, raw) -> str:
        if isinstance(raw, str):
            return raw.strip()

        if isinstance(raw, dict):
            return raw.get("generated_text", "").strip() or raw.get("text", "").strip()

        if isinstance(raw, list):
            if isinstance(raw[0], dict) and "generated_text" in raw[0]:
                return raw[0]["generated_text"].strip()
            return " ".join(str(x).strip() for x in raw if isinstance(x, str))

        return str(raw).strip()

    async def advanced_fusion(self, responses: Dict[str, str], query: str) -> str:
        cleaned = {m: self.clean_response(r) for m, r in responses.items()}
        seen = set()
        unique: List[Tuple[str, float, str]] = []

        for model, text in cleaned.items():
            if not text or any(self.is_similar(text, prev) for prev in seen):
                continue
            seen.add(text)
            score = self.score_response(text, query)
            unique.append((text, score, model))

        if not unique:
            return "⚠️ No high-quality content to merge."

        unique.sort(key=lambda x: x[1], reverse=True)
        best_responses = [await self.response_optimizer.optimize_response(t, query) for t, _, _ in unique[:3]]
        return "\n---\n".join(best_responses).strip()

    def is_similar(self, a: str, b: str, threshold: float = 0.9) -> bool:
        return SequenceMatcher(None, a.strip().lower(), b.strip().lower()).ratio() > threshold

    def score_response(self, response: str, query: str) -> float:
        try:
            tfidf = TfidfVectorizer().fit_transform([response, query])
            return cosine_similarity(tfidf[0:1], tfidf[1:2])[0][0]
        except:
            return 0.5