"""
请求粒度的上下文管理器

用于存储查询分析结果、各检索阶段中间结果、性能指标等。
支持线程安全的并发请求处理。
"""

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

from request_log_context import bind_request_log_context, reset_request_log_context


class RequestContextStage(Enum):
    """搜索阶段枚举"""
    TOTAL = "total_search"
    QUERY_PARSING = "query_parsing"
    BOOLEAN_PARSING = "boolean_parsing"
    QUERY_BUILDING = "query_building"
    # ES 主召回查询
    ELASTICSEARCH_SEARCH_PRIMARY = "elasticsearch_search_primary"
    # ES 按 ID 回源分页详情回填
    ELASTICSEARCH_PAGE_FILL = "elasticsearch_page_fill"
    RESULT_PROCESSING = "result_processing"
    COARSE_RANKING = "coarse_ranking"
    FINE_RANKING = "fine_ranking"
    RERANKING = "reranking"
    # 款式意图 SKU 预筛选（StyleSkuSelector.prepare_hits）
    STYLE_SKU_PREPARE_HITS = "style_sku_prepare_hits"


@dataclass
class QueryAnalysisResult:
    """查询分析结果"""
    original_query: Optional[str] = None
    query_normalized: Optional[str] = None
    rewritten_query: Optional[str] = None
    detected_language: Optional[str] = None
    translations: Dict[str, str] = field(default_factory=dict)
    keywords_queries: Dict[str, str] = field(default_factory=dict)
    query_vector: Optional[List[float]] = None
    boolean_ast: Optional[str] = None


@dataclass
class PerformanceMetrics:
    """性能指标"""
    stage_timings: Dict[str, float] = field(default_factory=dict)
    stage_start_times: Dict[str, float] = field(default_factory=dict)
    # 各阶段起止时间（Unix 毫秒，与 time.time() 一致）
    stage_time_bounds_ms: Dict[str, Dict[str, float]] = field(default_factory=dict)
    total_duration: float = 0.0
    extra_metrics: Dict[str, Any] = field(default_factory=dict)


class RequestContext:
    """
    请求粒度的上下文管理器

    功能：
    1. 存储查询分析结果和各阶段中间结果
    2. 自动跟踪各阶段耗时
    3. 提供线程安全的上下文访问
    4. 支持上下文管理器模式
    """

    def __init__(self, reqid: str = None, uid: str = None):
        # 生成唯一请求ID；如果外部未提供，则自动生成
        # 如果无法获取到 uid，则使用 "-1" 作为占位，用于日志关联
        self.reqid = reqid or str(uuid.uuid4())[:8]
        self.uid = uid or "-1"

        # 查询分析结果
        self.query_analysis = QueryAnalysisResult()

        # 各检索阶段中间结果
        self.intermediate_results = {
            'parsed_query': None,
            'query_node': None,
            'es_query': {},
            'es_response': {},
            'processed_hits': [],
            'raw_hits': []
        }

        # 性能指标
        self.performance_metrics = PerformanceMetrics()

        # 元数据
        self.metadata = {
            'search_params': {},  # size, from_, filters等
            'feature_flags': {},  # translation_enabled (由tenant_config控制), enable_embedding等
            'config_info': {},   # 索引配置、字段映射等
            'error_info': None,
            'warnings': []
        }

        # 日志记录器引用（延迟初始化）
        self._logger = None

    @property
    def logger(self):
        """获取日志记录器"""
        if self._logger is None:
            from utils.logger import get_logger
            self._logger = get_logger("request_context")
        return self._logger

    def start_stage(self, stage: RequestContextStage) -> float:
        """
        开始一个阶段的计时

        Args:
            stage: 阶段枚举

        Returns:
            开始时间戳
        """
        start_time = time.time()
        self.performance_metrics.stage_start_times[stage.value] = start_time
        self.logger.debug(
            f"Start stage | {stage.value}",
            extra={'reqid': self.reqid, 'uid': self.uid}
        )
        return start_time

    def end_stage(self, stage: RequestContextStage) -> float:
        """
        结束一个阶段的计时

        Args:
            stage: 阶段枚举

        Returns:
            阶段耗时（毫秒）
        """
        if stage.value not in self.performance_metrics.stage_start_times:
            self.logger.warning(
                f"Stage not started | {stage.value}",
                extra={'reqid': self.reqid, 'uid': self.uid}
            )
            return 0.0

        start_time = self.performance_metrics.stage_start_times[stage.value]
        end_time = time.time()
        duration_ms = (end_time - start_time) * 1000
        self.performance_metrics.stage_timings[stage.value] = duration_ms
        self.performance_metrics.stage_time_bounds_ms[stage.value] = {
            "start_unix_ms": round(start_time * 1000, 3),
            "end_unix_ms": round(end_time * 1000, 3),
        }

        self.logger.debug(
            f"End stage | {stage.value} | duration: {duration_ms:.2f}ms",
            extra={'reqid': self.reqid, 'uid': self.uid}
        )
        return duration_ms

    def get_stage_duration(self, stage: RequestContextStage) -> float:
        """
        获取指定阶段的耗时

        Args:
            stage: 阶段枚举

        Returns:
            阶段耗时（毫秒），如果未计时则返回0
        """
        return self.performance_metrics.stage_timings.get(stage.value, 0.0)

    def store_query_analysis(self, **kwargs) -> None:
        """
        存储查询分析结果

        Args:
            **kwargs: 查询分析相关的字段
        """
        for key, value in kwargs.items():
            if hasattr(self.query_analysis, key):
                setattr(self.query_analysis, key, value)
            else:
                self.logger.warning(
                    f"Unknown query analysis field | {key}",
                    extra={'reqid': self.reqid, 'uid': self.uid}
                )

    def store_intermediate_result(self, key: str, value: Any) -> None:
        """
        存储中间结果

        Args:
            key: 结果键名
            value: 结果值
        """
        self.intermediate_results[key] = value
        self.logger.debug(
            f"Store intermediate result | {key}",
            extra={'reqid': self.reqid, 'uid': self.uid}
        )

    def get_intermediate_result(self, key: str, default: Any = None) -> Any:
        """
        获取中间结果

        Args:
            key: 结果键名
            default: 默认值

        Returns:
            中间结果值
        """
        return self.intermediate_results.get(key, default)

    def add_warning(self, warning: str) -> None:
        """
        添加警告信息

        Args:
            warning: 警告信息
        """
        self.metadata['warnings'].append(warning)
        self.logger.warning(warning, extra={'reqid': self.reqid, 'uid': self.uid})

    def set_error(self, error: Exception) -> None:
        """
        设置错误信息

        Args:
            error: 异常对象
        """
        self.metadata['error_info'] = {
            'type': type(error).__name__,
            'message': str(error),
            'details': {}
        }
        self.logger.error(
            f"Set error info | {type(error).__name__}: {str(error)}",
            extra={'reqid': self.reqid, 'uid': self.uid}
        )

    def has_error(self) -> bool:
        """检查是否有错误"""
        return self.metadata['error_info'] is not None

    def calculate_stage_percentages(self) -> Dict[str, float]:
        """
        计算各阶段耗时占总耗时的百分比

        Returns:
            各阶段耗时占比字典
        """
        total = self.performance_metrics.total_duration
        if total <= 0:
            return {}

        percentages = {}
        for stage, duration in self.performance_metrics.stage_timings.items():
            percentages[stage] = round((duration / total) * 100, 2)

        return percentages

    def get_summary(self) -> Dict[str, Any]:
        """
        获取完整的上下文摘要

        Returns:
            包含所有关键信息的字典
        """
        return {
            'request_info': {
                'reqid': self.reqid,
                'uid': self.uid,
                'has_error': self.has_error(),
                'warnings_count': len(self.metadata['warnings'])
            },
            'query_analysis': {
                'original_query': self.query_analysis.original_query,
                'query_normalized': self.query_analysis.query_normalized,
                'rewritten_query': self.query_analysis.rewritten_query,
                'detected_language': self.query_analysis.detected_language,
            },
            'performance': {
                'total_duration_ms': round(self.performance_metrics.total_duration, 2),
                'stage_timings_ms': {
                    k: round(v, 2) for k, v in self.performance_metrics.stage_timings.items()
                },
                'stage_time_bounds_ms': {
                    stage: {
                        kk: round(vv, 3) for kk, vv in bounds.items()
                    }
                    for stage, bounds in self.performance_metrics.stage_time_bounds_ms.items()
                },
                'stage_percentages': self.calculate_stage_percentages()
            },
            'results': {
                'total_hits': len(self.intermediate_results.get('processed_hits', [])),
                'has_es_response': bool(self.intermediate_results.get('es_response')),
                'es_query_size': len(str(self.intermediate_results.get('es_query', {})))
            },
            'metadata': {
                'feature_flags': self.metadata['feature_flags'],
                'search_params': self.metadata['search_params'],
                'config_info': self.metadata['config_info']
            }
        }

    def log_performance_summary(self) -> None:
        """记录完整的性能摘要日志"""
        summary = self.get_summary()

        # 构建详细的日志消息
        msg_parts = [
            f"Search request performance summary | reqid: {self.reqid}",
            f"Total duration: {summary['performance']['total_duration_ms']:.2f}ms"
        ]

        # 添加各阶段耗时
        if summary['performance']['stage_timings_ms']:
            msg_parts.append("Stage durations:")
            bounds_map = summary['performance'].get('stage_time_bounds_ms') or {}
            for stage, duration in summary['performance']['stage_timings_ms'].items():
                percentage = summary['performance']['stage_percentages'].get(stage, 0)
                b = bounds_map.get(stage) or {}
                if b.get('start_unix_ms') is not None and b.get('end_unix_ms') is not None:
                    msg_parts.append(
                        f"  - {stage}: {duration:.2f}ms ({percentage}%) "
                        f"[{b['start_unix_ms']:.3f} -> {b['end_unix_ms']:.3f} ms]"
                    )
                else:
                    msg_parts.append(f"  - {stage}: {duration:.2f}ms ({percentage}%)")

        # 添加查询信息
        if summary['query_analysis']['original_query']:
            msg_parts.append(
                "Query: "
                f"'{summary['query_analysis']['original_query']}' "
                f"-> '{summary['query_analysis']['rewritten_query']}' "
                f"({summary['query_analysis']['detected_language']})"
            )

        # 添加结果统计
        msg_parts.append(
            f"Results: {summary['results']['total_hits']} hits "
            f"ES query size: {summary['results']['es_query_size']} chars"
        )

        # 添加错误信息（如果有）
        if summary['request_info']['has_error']:
            error_info = self.metadata['error_info']
            msg_parts.append(f"Error: {error_info['type']}: {error_info['message']}")

        # 添加警告信息（如果有）
        if summary['request_info']['warnings_count'] > 0:
            msg_parts.append(f"Warnings: {summary['request_info']['warnings_count']}")

        log_message = " | ".join(msg_parts)

        if self.has_error():
            self.logger.error(log_message, extra={'extra_data': summary, 'reqid': self.reqid, 'uid': self.uid})
        elif summary['request_info']['warnings_count'] > 0:
            self.logger.warning(log_message, extra={'extra_data': summary, 'reqid': self.reqid, 'uid': self.uid})
        else:
            self.logger.info(log_message, extra={'extra_data': summary, 'reqid': self.reqid, 'uid': self.uid})

    def __enter__(self):
        """上下文管理器入口"""
        self.start_stage(RequestContextStage.TOTAL)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        # 结束总计时
        self.end_stage(RequestContextStage.TOTAL)
        self.performance_metrics.total_duration = self.get_stage_duration(RequestContextStage.TOTAL)

        # 记录性能摘要
        self.log_performance_summary()

        # 如果有异常，记录错误信息
        if exc_type and exc_val:
            self.set_error(exc_val)


# 便利函数
def create_request_context(reqid: str = None, uid: str = None) -> RequestContext:
    """创建新的请求上下文"""
    return RequestContext(reqid, uid)


def get_current_request_context() -> Optional[RequestContext]:
    """获取当前线程的请求上下文（如果已设置）"""
    return getattr(threading.current_thread(), 'request_context', None)


def set_current_request_context(context: RequestContext) -> None:
    """设置当前线程的请求上下文"""
    threading.current_thread().request_context = context
    _, _, tokens = bind_request_log_context(context.reqid, context.uid)
    threading.current_thread().request_log_tokens = tokens


def clear_current_request_context() -> None:
    """清除当前线程的请求上下文"""
    tokens = getattr(threading.current_thread(), 'request_log_tokens', None)
    if tokens is not None:
        reset_request_log_context(tokens)
        delattr(threading.current_thread(), 'request_log_tokens')
    if hasattr(threading.current_thread(), 'request_context'):
        delattr(threading.current_thread(), 'request_context')