src/integration/query_handler.py

import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

from agent.intent_classifier import IntentType
from retrieval import get_retrieval_engine
from llm_generator import get_llm_generator, DEFAULT_LLM_MODEL
from config import ENABLE_QUERY_REWRITING, CONVERSATION_SUMMARIZATION_ENABLED, CONVERSATION_MAX_TURNS_BEFORE_SUMMARY, DATABASE_CONFIG
from query_rewriter import get_query_rewriter
from utils import get_conversation_summarizer, get_metrics, record_query
from integration.context_manager import get_context_manager
from integration.intent_router import get_intent_router, RoutingDecision
from integration.entity_context import get_entity_store, FollowUpDetection
from database.view_manager import extract_identifier, extract_credit_file_id, has_database_signals, query_and_format, IdentifierType, check_dealer_access, _DEALER_ACCESS_DENIED_MSG

logger = logging.getLogger(__name__)

_CRM_BLOCKED_KEYWORDS = [
    "tbo", "cbo",
    "toplam bor\u00e7", "toplam borc",
    "bor\u00e7 oran\u0131", "borc orani",
    "bor\u00e7 y\u00fck\u00fc", "borc yuku",
    "risk", "risk skoru", "risk durumu",
    "risk analizi", "risk de\u011ferlendirmesi",
    "kredi riski", "risk puan\u0131", "risk puani",
    "finansal risk", "bor\u00e7luluk", "borcululuk",
]

_CRM_REFERRAL_RESPONSE = (
    "Bu konu (TBO, CBO, toplam bor\u00e7 veya risk de\u011ferlendirmesi), "
    "CRM ekibinin sorumluluk alan\u0131na girmektedir.\n"
    "Daha fazla bilgi ve destek i\u00e7in l\u00fctfen CRM ekibiyle ileti\u015fime ge\u00e7iniz."
)


def _is_crm_blocked(query: str) -> bool:
    q_lower = query.lower()
    return any(kw in q_lower for kw in _CRM_BLOCKED_KEYWORDS)


class QueryMode(Enum):
    DOCUMENTS = "documents"
    DATABASE = "database"
    HYBRID = "hybrid"
    AUTO = "auto"
    CONVERSATIONAL = "conversational"


@dataclass
class QueryResult:
    query: str
    response: str
    mode: QueryMode
    sources: List[Dict[str, Any]] = field(default_factory=list)
    execution_time: float = 0.0
    intent: Optional[str] = None
    sql_query: Optional[str] = None
    db_results: Optional[List[Dict[str, Any]]] = None
    error: Optional[str] = None
    rewrite_metadata: Optional[Dict[str, Any]] = None
    summarization_metadata: Optional[Dict[str, Any]] = None
    context_filtering_metadata: Optional[Dict[str, Any]] = None
    routing_metadata: Optional[Dict[str, Any]] = None
    model_id: Optional[str] = None

    @property
    def success(self) -> bool:
        return self.error is None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "query": self.query,
            "response": self.response,
            "mode": self.mode.value,
            "sources": self.sources,
            "execution_time": self.execution_time,
            "intent": self.intent,
            "sql_query": self.sql_query,
            "db_results": self.db_results,
            "error": self.error,
            "success": self.success,
            "rewrite_metadata": self.rewrite_metadata,
            "summarization_metadata": self.summarization_metadata,
            "context_filtering_metadata": self.context_filtering_metadata,
            "routing_metadata": self.routing_metadata,
            "model_id": self.model_id,
        }


class QueryHandler:
    def __init__(self, default_mode: QueryMode = QueryMode.AUTO, **_kwargs):
        self._default_mode = default_mode
        self._query_history: List[QueryResult] = []

    @property
    def default_mode(self) -> QueryMode:
        return self._default_mode

    @default_mode.setter
    def default_mode(self, mode: QueryMode) -> None:
        self._default_mode = mode

    @property
    def query_history(self) -> List[QueryResult]:
        return self._query_history.copy()

    def clear_history(self) -> None:
        self._query_history.clear()

    def execute(
        self,
        query: str,
        mode: Optional[QueryMode] = None,
        chat_history: Optional[List[Dict[str, str]]] = None,
        model_id: str = DEFAULT_LLM_MODEL,
    ) -> QueryResult:
        start_time = time.time()
        effective_mode = mode or self._default_mode
        original_query = query

        if _is_crm_blocked(query):
            execution_time = time.time() - start_time
            result = QueryResult(
                query=original_query,
                response=_CRM_REFERRAL_RESPONSE,
                mode=effective_mode,
                execution_time=execution_time,
                intent="crm_referral",
            )
            self._query_history.append(result)
            self._record_metrics(result)
            return result

        entity_store = get_entity_store()
        entity_store.advance_turn()

        chat_history = self._prepare_chat_history(chat_history)
        resolved_query, identifier_type, identifier_value, rewrite_metadata, is_followup = self._resolve_query(query, chat_history)

        if identifier_type == IdentifierType.CREDIT_FILE_ID and identifier_value:
            if not check_dealer_access(identifier_value):
                execution_time = time.time() - start_time
                result = QueryResult(
                    query=original_query,
                    response=_DEALER_ACCESS_DENIED_MSG,
                    mode=effective_mode,
                    execution_time=execution_time,
                    intent="dealer_access_denied",
                    error="dealer_access_denied",
                )
                self._query_history.append(result)
                self._record_metrics(result)
                return result

        followup_detection = getattr(self, "_last_followup_detection", None)
        if followup_detection and followup_detection.needs_clarification:
            execution_time = time.time() - start_time
            result = QueryResult(
                query=original_query,
                response="Hangi teklif, dosya veya musteri hakkinda bilgi istediginizi belirtir misiniz? Ornegin teklif numarasi veya musteri referansi ile tekrar sorabilirsiniz.",
                mode=effective_mode,
                execution_time=execution_time,
                intent="clarification",
                routing_metadata={"followup_detection": followup_detection.to_dict()},
            )
            self._query_history.append(result)
            self._record_metrics(result)
            return result

        try:
            if effective_mode == QueryMode.DATABASE:
                result = self._dispatch_database(resolved_query, start_time, original_query, identifier_type, identifier_value)
            elif effective_mode == QueryMode.DOCUMENTS:
                result = self._execute_documents(resolved_query, chat_history, start_time, original_query, model_id)
            elif effective_mode == QueryMode.HYBRID:
                result = self._execute_hybrid(resolved_query, chat_history, start_time, original_query, identifier_type, identifier_value, model_id, is_followup)
            else:
                result = self._dispatch_auto(resolved_query, chat_history, start_time, original_query, identifier_type, identifier_value, model_id, is_followup)

            result.rewrite_metadata = rewrite_metadata
            result.summarization_metadata = getattr(self, "_last_summarization_metadata", None)
            result.model_id = model_id
            entity_store.last_query_mode = result.mode.value
            self._record_metrics(result)
            return result
        except Exception as e:
            execution_time = time.time() - start_time
            result = QueryResult(
                query=original_query,
                response=f"Sorgu islenirken hata olustu: {str(e)}",
                mode=effective_mode,
                execution_time=execution_time,
                error=str(e),
                rewrite_metadata=rewrite_metadata,
                summarization_metadata=getattr(self, "_last_summarization_metadata", None),
            )
            self._record_metrics(result)
            self._query_history.append(result)
            return result

    def _dispatch_database(self, query, start_time, original_query, identifier_type, identifier_value):
        database_available = DATABASE_CONFIG.get("enabled", False)
        if not database_available:
            return self._execute_documents(query, [], start_time, original_query)
        return self._execute_database(query, start_time, original_query, identifier_type, identifier_value)

    def _dispatch_auto(self, query, chat_history, start_time, original_query, identifier_type, identifier_value, model_id, is_followup):
        database_available = DATABASE_CONFIG.get("enabled", False)

        if is_followup and identifier_type and identifier_value and database_available:
            return self._execute_database(query, start_time, original_query, identifier_type, identifier_value)

        router = get_intent_router()
        decision = router.route(original_query, database_available=database_available)

        if decision.execution_mode == "database":
            return self._execute_database(query, start_time, original_query, identifier_type, identifier_value)
        if decision.execution_mode == "conversational":
            return self._execute_conversational(query, chat_history, start_time, original_query, model_id)
        return self._execute_documents(query, chat_history, start_time, original_query, model_id)

    def _execute_hybrid(
        self,
        query: str,
        chat_history: Optional[List[Dict[str, str]]],
        start_time: float,
        original_query: str,
        identifier_type: Optional[IdentifierType] = None,
        identifier_value: Optional[str] = None,
        model_id: str = DEFAULT_LLM_MODEL,
        is_followup: bool = False,
    ) -> QueryResult:
        database_available = DATABASE_CONFIG.get("enabled", False)
        db_result_data = None
        doc_result = None

        should_try_db = has_database_signals(original_query, id_type=identifier_type, id_value=identifier_value) or (is_followup and identifier_type and identifier_value)
        if database_available and should_try_db:
            db_result_data = query_and_format(original_query, identifier_type=identifier_type, identifier_value=identifier_value)

        try:
            engine = get_retrieval_engine()
            generator = get_llm_generator()
            results, _ = engine.retrieve(query, chat_history=chat_history, use_reranking=True, skip_rewrite=True)
            context = engine.build_context(results)
            if results:
                doc_response = generator.generate(query, context, chat_history, model_id)
                doc_result = {"response": doc_response, "sources": results}
        except Exception as e:
            logger.warning(f"Hybrid doc retrieval failed: {e}")

        if db_result_data and db_result_data.get("success") and db_result_data.get("rows"):
            execution_time = time.time() - start_time
            ctx = get_context_manager()
            ctx.set_last_database_response(db_result_data["response"])
            entity_store = get_entity_store()
            if identifier_type and identifier_value:
                entity_store.update_entity(
                    identifier_type=identifier_type,
                    identifier_value=identifier_value,
                    data_payload=db_result_data.get("rows"),
                    response_text=db_result_data["response"],
                    sql_query=db_result_data.get("sql"),
                )
            elif is_followup:
                entity_store.update_turn_only()
            result = QueryResult(
                query=original_query,
                response=db_result_data["response"],
                mode=QueryMode.HYBRID,
                execution_time=execution_time,
                intent=IntentType.DATABASE_QUERY.value,
                sql_query=db_result_data.get("sql"),
                db_results=db_result_data.get("rows"),
                routing_metadata={"hybrid_source": "database", "doc_available": doc_result is not None},
            )
            self._query_history.append(result)
            return result

        if doc_result:
            execution_time = time.time() - start_time
            sources = [
                {
                    "content": r.get("text", "")[:200] + "..." if len(r.get("text", "")) > 200 else r.get("text", ""),
                    "metadata": r.get("metadata", {}),
                    "source": r.get("source", ""),
                    "score": r.get("score", 0.0),
                }
                for r in doc_result["sources"]
            ]
            result = QueryResult(
                query=original_query,
                response=doc_result["response"],
                mode=QueryMode.HYBRID,
                sources=sources,
                execution_time=execution_time,
                intent=IntentType.DOCUMENT_QUERY.value,
                routing_metadata={"hybrid_source": "documents", "db_attempted": db_result_data is not None},
            )
            self._query_history.append(result)
            return result

        if db_result_data and db_result_data.get("success"):
            execution_time = time.time() - start_time
            result = QueryResult(
                query=original_query,
                response=db_result_data["response"],
                mode=QueryMode.HYBRID,
                execution_time=execution_time,
                intent=IntentType.DATABASE_QUERY.value,
                sql_query=db_result_data.get("sql"),
                db_results=db_result_data.get("rows", []),
                routing_metadata={"hybrid_source": "database_no_rows"},
            )
            self._query_history.append(result)
            return result

        execution_time = time.time() - start_time
        result = QueryResult(
            query=original_query,
            response="Hem veritabani hem de dokuman aramasinda sonuc bulunamadi.",
            mode=QueryMode.HYBRID,
            execution_time=execution_time,
            error="no_results_from_either_source",
            routing_metadata={"hybrid_source": "none"},
        )
        self._query_history.append(result)
        return result

    def _execute_database(
        self,
        query: str,
        start_time: float,
        original_query: str,
        identifier_type: Optional[IdentifierType] = None,
        identifier_value: Optional[str] = None,
    ) -> QueryResult:
        db_result = query_and_format(original_query, identifier_type=identifier_type, identifier_value=identifier_value)
        execution_time = time.time() - start_time

        if db_result["success"]:
            ctx = get_context_manager()
            ctx.set_last_database_response(db_result["response"])
            entity_store = get_entity_store()
            if identifier_type and identifier_value:
                entity_store.update_entity(
                    identifier_type=identifier_type,
                    identifier_value=identifier_value,
                    data_payload=db_result.get("rows"),
                    response_text=db_result["response"],
                    sql_query=db_result.get("sql"),
                )
            elif db_result.get("is_aggregate") and db_result.get("view_name"):
                entity_store.update_aggregate(
                    view_name=db_result["view_name"],
                    response_text=db_result["response"],
                    sql_query=db_result.get("sql"),
                )

        result = QueryResult(
            query=original_query,
            response=db_result["response"],
            mode=QueryMode.DATABASE,
            execution_time=execution_time,
            intent=IntentType.DATABASE_QUERY.value,
            sql_query=db_result.get("sql"),
            db_results=db_result.get("rows"),
            error=db_result.get("error") if not db_result["success"] else None,
        )
        self._query_history.append(result)
        return result

    def _execute_documents(
        self,
        query: str,
        chat_history: Optional[List[Dict[str, str]]],
        start_time: float,
        original_query: str,
        model_id: str = DEFAULT_LLM_MODEL,
    ) -> QueryResult:
        try:
            engine = get_retrieval_engine()
            generator = get_llm_generator()
            results, _ = engine.retrieve(query, chat_history=chat_history, use_reranking=True, skip_rewrite=True)
            context = engine.build_context(results)
            preamble = get_entity_store().build_context_preamble()
            if preamble:
                context = f"{preamble}\n\n{context}"
            response = generator.generate(query, context, chat_history, model_id)
            execution_time = time.time() - start_time

            sources = [
                {
                    "content": r.get("text", "")[:200] + "..." if len(r.get("text", "")) > 200 else r.get("text", ""),
                    "metadata": r.get("metadata", {}),
                    "source": r.get("source", ""),
                    "score": r.get("score", 0.0),
                }
                for r in results
            ]

            result = QueryResult(
                query=original_query,
                response=response,
                mode=QueryMode.DOCUMENTS,
                sources=sources,
                execution_time=execution_time,
                intent=IntentType.DOCUMENT_QUERY.value,
            )
            self._query_history.append(result)
            return result
        except Exception as e:
            execution_time = time.time() - start_time
            result = QueryResult(
                query=original_query,
                response=f"Dokuman arama hatasi: {str(e)}",
                mode=QueryMode.DOCUMENTS,
                execution_time=execution_time,
                error=str(e),
            )
            self._query_history.append(result)
            return result

    def _execute_conversational(
        self,
        query: str,
        chat_history: Optional[List[Dict[str, str]]],
        start_time: float,
        original_query: str,
        model_id: str = DEFAULT_LLM_MODEL,
    ) -> QueryResult:
        try:
            generator = get_llm_generator()
            context = "Bu bir sohbet mesajidir. Kullaniciyla nazik ve yardimci bir sekilde konusun."
            preamble = get_entity_store().build_context_preamble()
            if preamble:
                context = f"{preamble}\n\n{context}"
            response = generator.generate(query, context, chat_history, model_id)
            execution_time = time.time() - start_time

            result = QueryResult(
                query=original_query,
                response=response,
                mode=QueryMode.CONVERSATIONAL,
                execution_time=execution_time,
                intent=IntentType.CONVERSATIONAL.value,
            )
            self._query_history.append(result)
            return result
        except Exception as e:
            execution_time = time.time() - start_time
            result = QueryResult(
                query=original_query,
                response=f"Yanit olusturulamadi: {str(e)}",
                mode=QueryMode.CONVERSATIONAL,
                execution_time=execution_time,
                error=str(e),
            )
            self._query_history.append(result)
            return result

    def _prepare_chat_history(self, chat_history: Optional[List[Dict[str, str]]]) -> List[Dict[str, str]]:
        if not chat_history:
            self._last_summarization_metadata = {"conversation_summarized": False, "original_message_count": 0, "summarized_message_count": 0}
            return []

        cleaned = [msg for msg in chat_history if isinstance(msg, dict) and "role" in msg and "content" in msg]
        original_count = len(cleaned)
        summarized = False
        summarization_time_ms = 0.0

        if CONVERSATION_SUMMARIZATION_ENABLED and len(cleaned) > CONVERSATION_MAX_TURNS_BEFORE_SUMMARY * 2:
            try:
                t0 = time.time()
                summarizer = get_conversation_summarizer()
                cleaned = summarizer.summarize_if_needed(cleaned)
                summarization_time_ms = (time.time() - t0) * 1000
                summarized = len(cleaned) != original_count
            except Exception as e:
                logger.warning(f"Conversation summarization failed: {e}")

        self._last_summarization_metadata = {
            "conversation_summarized": summarized,
            "original_message_count": original_count,
            "summarized_message_count": len(cleaned),
            "summarization_time_ms": summarization_time_ms,
        }
        return cleaned

    def _resolve_query(self, query: str, chat_history: Optional[List[Dict[str, str]]]):
        resolved_text = query
        rewrite_metadata = None

        id_type, id_value = extract_identifier(query)

        is_followup = False
        self._last_followup_detection = None
        if not id_value:
            entity_store = get_entity_store()
            detection = entity_store.detect_followup(query)
            self._last_followup_detection = detection
            is_followup = detection.is_followup
            if is_followup and detection.resolved_identifier_type:
                id_type = detection.resolved_identifier_type
                id_value = detection.resolved_identifier_value
                logger.info(
                    "followup resolved %s=%s for query: %s",
                    id_type.value if id_type else "none",
                    id_value,
                    query[:80],
                )
            if is_followup and not detection.needs_clarification and id_value:
                entity_store.update_turn_only()

        if not is_followup and ENABLE_QUERY_REWRITING:
            try:
                rewriter = get_query_rewriter()
                resolved_text, metadata = rewriter.rewrite(query, chat_history or [])
                rewrite_metadata = metadata
                if resolved_text != query:
                    logger.info(f"Query rewritten: '{query[:80]}' -> '{resolved_text[:80]}'")
            except Exception as e:
                logger.warning(f"Query rewriting failed: {e}")
                rewrite_metadata = {"method": "fallback_failed", "error": str(e)}
                resolved_text = query

        if id_type and id_value:
            try:
                ctx = get_context_manager()
                ctx.set_active_identifier(id_type, id_value)
            except Exception:
                pass

        return resolved_text, id_type, id_value, rewrite_metadata, is_followup

    def _record_metrics(self, result: QueryResult) -> None:
        try:
            metrics = get_metrics()
            metrics.record(query=result.query, mode=result.mode.value, response_time=result.execution_time, success=result.success, error=result.error)
        except Exception:
            pass


_query_handler_instance: Optional[QueryHandler] = None


def get_query_handler(default_mode: QueryMode = QueryMode.AUTO, **_kwargs) -> QueryHandler:
    global _query_handler_instance
    if _query_handler_instance is None:
        _query_handler_instance = QueryHandler(default_mode=default_mode)
    return _query_handler_instance


def reset_query_handler() -> None:
    global _query_handler_instance
    _query_handler_instance = None