src/query_rewriter.py

import re
import time
import hashlib
import logging
from typing import List, Dict, Optional, Tuple
from functools import lru_cache
from groq import Groq

from config import (
    ENABLE_QUERY_REWRITING,
    MIN_HISTORY_FOR_REWRITE,
    MAX_HISTORY_TURNS_FOR_REWRITE,
    MAX_CONTENT_LENGTH_FOR_REWRITE,
    QUERY_REWRITE_MODEL,
    QUERY_REWRITE_TEMPERATURE,
    QUERY_REWRITE_MAX_TOKENS,
    QUERY_REWRITE_TIMEOUT,
    GROQ_API_KEY,
    QUERY_REWRITE_CACHE_SIZE,
    QUERY_REWRITE_CACHE_TTL_SECONDS
)
from prompts import QUERY_REWRITE_PROMPT

logger = logging.getLogger(__name__)

_FOLLOW_UP_INDICATORS = frozenset([
    "bu", "bunu", "bunun", "buna", "bunlar",
    "o", "onu", "onun", "ona", "onlar",
    "su", "sunu", "sunun",
    "boyle", "oyle", "soyle",
    "burada", "orada", "surada",
    "burdaki", "ordaki",
    "ayni", "diger", "sonraki", "onceki",
    "bahsettigin", "dedin", "soyledin",
    "yukardaki", "asagidaki",
])

_PRONOUN_REPLACEMENTS = {
    "bu": "{topic}",
    "bunu": "{topic}'nu",
    "bunun": "{topic}'nun",
    "buna": "{topic}'na",
    "bunlar": "{topic} ve ilgili konular",
    "o": "{topic}",
    "onu": "{topic}'nu",
    "onun": "{topic}'nun",
    "ona": "{topic}'na",
    "su": "{topic}",
    "sunu": "{topic}'nu",
    "sunun": "{topic}'nun",
}


class QueryRewriter:
    
    def __init__(self, cache_duration_seconds: int = QUERY_REWRITE_CACHE_TTL_SECONDS):
        self._indicators = _FOLLOW_UP_INDICATORS
        self._pattern = self._build_pattern()
        self._client = None
        self._cache: Dict[str, str] = {}
        self._cache_ttl: Dict[str, float] = {}
        self._max_cache_size: int = QUERY_REWRITE_CACHE_SIZE
        self._cache_duration: int = cache_duration_seconds
        self._cache_hits: int = 0
        self._cache_misses: int = 0
    
    @property
    def client(self) -> Groq:
        if self._client is None:
            self._client = Groq(api_key=GROQ_API_KEY, timeout=QUERY_REWRITE_TIMEOUT)
        return self._client
    
    def _build_pattern(self) -> re.Pattern:
        escaped = [re.escape(w) for w in self._indicators]
        return re.compile(r'\b(' + '|'.join(escaped) + r')\b', re.IGNORECASE)
    
    def is_follow_up(self, query: str, chat_history: List[Dict[str, str]]) -> bool:
        if not ENABLE_QUERY_REWRITING:
            return False
        
        if len(chat_history) < MIN_HISTORY_FOR_REWRITE * 2:
            return False
        
        query_lower = query.lower().strip()
        words = query_lower.split()
        
        subject_indicators = ['atlas', 'erp', 'sistem', 'modul', 'program', 'uygulama', 'yazılım', 'şirket', 'firma']
        has_explicit_subject = any(subj in query_lower for subj in subject_indicators)
        
        if self._pattern.search(query_lower):
            return True
        
        if query_lower.startswith(('peki', 'ya ', 'ee', 'hmm', 'tamam', 'iyi', 'anladım')):
            return True
        
        if has_explicit_subject:
            return False
        
        if len(words) <= 5:
            return True
        
        question_starters = ['ne', 'nasıl', 'neden', 'niçin', 'kim', 'nerede', 'hangi', 'kaç', 'ne?', 'nasıl?']
        if words and words[0].rstrip('?') in [q.rstrip('?') for q in question_starters]:
            return True
        
        if len(words) <= 8 and query_lower.endswith('?'):
            return True
        
        return False
    
    def format_history(self, chat_history: List[Dict[str, str]]) -> str:
        if not chat_history:
            return ""
        
        recent = chat_history[-(MAX_HISTORY_TURNS_FOR_REWRITE * 2):]
        parts = []
        
        for msg in recent:
            role = "Kullanici" if msg["role"] == "user" else "Asistan"
            content = msg.get("content", "")[:MAX_CONTENT_LENGTH_FOR_REWRITE]
            parts.append(f"{role}: {content}")
        
        return "\n".join(parts)
    
    def get_last_topic(self, chat_history: List[Dict[str, str]]) -> Optional[str]:
        for msg in reversed(chat_history):
            if msg["role"] == "user":
                return msg.get("content", "")[:100]
        return None
    
    def _get_cache_key(self, query: str, chat_history: List[Dict[str, str]]) -> str:
        history_part = "|".join(m.get("content", "")[:50] for m in chat_history[-4:])
        raw = f"{query}::{history_part}"
        return hashlib.md5(raw.encode("utf-8")).hexdigest()

    def _get_from_cache(self, cache_key: str) -> Optional[str]:
        if cache_key not in self._cache:
            self._cache_misses += 1
            return None
        age = time.time() - self._cache_ttl.get(cache_key, 0)
        if age >= self._cache_duration:
            del self._cache[cache_key]
            del self._cache_ttl[cache_key]
            self._cache_misses += 1
            logger.debug(f"Cache entry expired after {age:.1f}s")
            return None
        self._cache_hits += 1
        return self._cache[cache_key]

    def _add_to_cache(self, cache_key: str, rewritten: str) -> None:
        if len(self._cache) >= self._max_cache_size and cache_key not in self._cache:
            oldest_key = min(self._cache_ttl, key=self._cache_ttl.get)
            del self._cache[oldest_key]
            del self._cache_ttl[oldest_key]
            logger.debug("Cache evicted oldest entry due to size limit")
        self._cache[cache_key] = rewritten
        self._cache_ttl[cache_key] = time.time()
    
    def rewrite(self, query: str, chat_history: List[Dict[str, str]]) -> Tuple[str, dict]:
        start_time = time.time()
        debug_info = {
            "original_query": query,
            "is_follow_up": False,
            "rewritten_query": query,
            "method": "none",
            "rewrite_time_ms": 0
        }
        
        if not self.is_follow_up(query, chat_history):
            debug_info["rewrite_time_ms"] = int((time.time() - start_time) * 1000)
            return query, debug_info
        
        debug_info["is_follow_up"] = True

        cache_key = self._get_cache_key(query, chat_history)
        cached = self._get_from_cache(cache_key)
        if cached is not None:
            debug_info["rewritten_query"] = cached
            debug_info["method"] = "cache"
            debug_info["cache_hit"] = True
            debug_info["cache_size"] = len(self._cache)
            debug_info["cache_age_seconds"] = round(time.time() - self._cache_ttl.get(cache_key, time.time()), 2)
            debug_info["rewrite_time_ms"] = int((time.time() - start_time) * 1000)
            return cached, debug_info

        rewritten = self._safe_rewrite(query, chat_history, debug_info)
        self._add_to_cache(cache_key, rewritten)
        debug_info["rewritten_query"] = rewritten
        debug_info["cache_hit"] = False
        debug_info["cache_size"] = len(self._cache)
        
        debug_info["rewrite_time_ms"] = int((time.time() - start_time) * 1000)
        return rewritten, debug_info
    
    def _safe_rewrite(self, query: str, chat_history: List[Dict[str, str]], debug_info: dict) -> str:
        try:
            rewritten = self._llm_rewrite(query, chat_history)
            debug_info["method"] = "llm"
            return rewritten
        except Exception as e:
            logger.warning(f"LLM rewrite failed: {e}")
            debug_info["llm_error"] = str(e)
        
        try:
            rewritten = self._rule_based_rewrite(query, chat_history)
            if rewritten != query:
                debug_info["method"] = "rule"
                return rewritten
        except Exception as e:
            logger.warning(f"Rule-based rewrite failed: {e}")
        
        try:
            rewritten = self._append_context_fallback(query, chat_history)
            if rewritten != query:
                debug_info["method"] = "append"
                return rewritten
        except Exception as e:
            logger.warning(f"Append context failed: {e}")
        
        debug_info["method"] = "fallback_failed"
        return query
    
    def _llm_rewrite(self, query: str, chat_history: List[Dict[str, str]]) -> str:
        history_str = self.format_history(chat_history)
        prompt = QUERY_REWRITE_PROMPT.format(
            chat_history=history_str,
            original_query=query
        )
        
        response = self.client.chat.completions.create(
            model=QUERY_REWRITE_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=QUERY_REWRITE_TEMPERATURE,
            max_tokens=QUERY_REWRITE_MAX_TOKENS
        )
        
        rewritten = response.choices[0].message.content.strip()
        rewritten = self._clean_response(rewritten)
        
        if not rewritten or len(rewritten) < 3:
            return query
        
        return rewritten
    
    def _clean_response(self, response: str) -> str:
        lines = response.split('\n')
        result = lines[0].strip()
        
        for prefix in ['soru:', 'yeniden yazilmis soru:', 'cevap:']:
            if result.lower().startswith(prefix):
                result = result[len(prefix):].strip()
        
        result = result.strip('"\'')
        return result
    
    def _rule_based_rewrite(self, query: str, chat_history: List[Dict[str, str]]) -> str:
        topic = self._extract_topic(chat_history)
        if not topic:
            return query
        
        result = query
        query_lower = query.lower()
        
        for pronoun, replacement in _PRONOUN_REPLACEMENTS.items():
            pattern = re.compile(r'\b' + re.escape(pronoun) + r'\b', re.IGNORECASE)
            if pattern.search(query_lower):
                filled = replacement.format(topic=topic)
                result = pattern.sub(filled, result, count=1)
                break
        
        return result
    
    def _append_context_fallback(self, query: str, chat_history: List[Dict[str, str]]) -> str:
        topic = self._extract_topic(chat_history)
        if not topic:
            return query
        return f"{topic} hakkinda: {query}"
    
    def _extract_topic(self, chat_history: List[Dict[str, str]]) -> Optional[str]:
        last_q = self.get_last_topic(chat_history)
        if not last_q:
            return None
        
        last_q = last_q.strip().rstrip('?')
        words = last_q.split()
        
        stop_words = {'nedir', 'nelerdir', 'nasıl', 'ne', 'kadar', 'mi', 'mu'}
        topic_words = [w for w in words if w.lower() not in stop_words]
        
        if topic_words:
            return ' '.join(topic_words[:4])
        return last_q[:50]
    
    def clear_cache(self):
        self._cache.clear()
        self._cache_ttl.clear()

    def get_cache_stats(self) -> Dict[str, any]:
        now = time.time()
        ages = [now - ts for ts in self._cache_ttl.values()] if self._cache_ttl else []
        total_requests = self._cache_hits + self._cache_misses
        return {
            "total_entries": len(self._cache),
            "max_size": self._max_cache_size,
            "ttl_seconds": self._cache_duration,
            "hit_rate": round(self._cache_hits / total_requests, 4) if total_requests > 0 else 0.0,
            "cache_hits": self._cache_hits,
            "cache_misses": self._cache_misses,
            "oldest_entry_age": round(max(ages), 2) if ages else 0.0,
            "newest_entry_age": round(min(ages), 2) if ages else 0.0,
        }

    def warm_cache(self, common_queries: List[Tuple[str, str]]) -> None:
        for query, rewrite in common_queries:
            cache_key = hashlib.md5(f"{query}::".encode("utf-8")).hexdigest()
            self._add_to_cache(cache_key, rewrite)

    def reset_metrics(self) -> None:
        self._cache_hits = 0
        self._cache_misses = 0


_rewriter_instance: Optional[QueryRewriter] = None


def get_query_rewriter() -> QueryRewriter:
    global _rewriter_instance
    if _rewriter_instance is None:
        _rewriter_instance = QueryRewriter()
    return _rewriter_instance


def reset_query_rewriter():
    global _rewriter_instance
    _rewriter_instance = None