"""Token counting functionality for LLM usage tracking."""

import inspect
import json
import time
from datetime import datetime, timedelta, UTC
from pathlib import Path
from typing import Any, Dict, List, Optional

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from loguru import logger
from sqlalchemy import func, text

from ..database.models import ModelUsage, TokenUsage
from .query_utils import get_research_mode_condition, get_time_filter_condition


class TokenCountingCallback(BaseCallbackHandler):
    """Callback handler for counting tokens across different models."""

    def __init__(
        self,
        research_id: Optional[str] = None,
        research_context: Optional[Dict[str, Any]] = None,
    ):
        """Initialize the token counting callback.

        Args:
            research_id: The ID of the research to track tokens for
            research_context: Additional research context for enhanced tracking
        """
        super().__init__()
        self.research_id = research_id
        self.research_context = research_context or {}
        self.current_model = None
        self.current_provider = None
        self.preset_model = None  # Model name set during callback creation
        self.preset_provider = None  # Provider set during callback creation

        # Phase 1 Enhancement: Track timing and context
        self.start_time = None
        self.response_time_ms = None
        self.success_status = "success"
        self.error_type = None

        # Call stack tracking
        self.calling_file = None
        self.calling_function = None
        self.call_stack = None

        # Context overflow tracking
        self.context_limit = None
        self.context_truncated = False
        self.tokens_truncated = 0
        self.truncation_ratio = 0.0
        self.original_prompt_estimate = 0

        # Raw Ollama response metrics
        self.ollama_metrics = {}

        # Track token counts in memory
        self.counts = {
            "total_tokens": 0,
            "total_prompt_tokens": 0,
            "total_completion_tokens": 0,
            "by_model": {},
        }

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Called when LLM starts running."""
        # Phase 1 Enhancement: Start timing
        self.start_time = time.time()

        # Estimate original prompt size (rough estimate: ~4 chars per token)
        if prompts:
            total_chars = sum(len(prompt) for prompt in prompts)
            self.original_prompt_estimate = total_chars // 4
            logger.debug(
                f"Estimated prompt tokens: {self.original_prompt_estimate} (from {total_chars} chars)"
            )

        # Get context limit from research context (will be set from settings)
        self.context_limit = self.research_context.get("context_limit")
        logger.debug(
            f"TokenCounter initialized - context_limit from research_context: {self.context_limit}, "
            f"research_context keys: {list(self.research_context.keys()) if self.research_context else 'None'}"
        )

        # Phase 1 Enhancement: Capture call stack information
        try:
            stack = inspect.stack()

            # Skip the first few frames (this method, langchain internals)
            # Look for the first frame that's in our project directory
            for frame_info in stack[1:]:
                file_path = frame_info.filename
                # Look for any frame containing local_deep_research project
                if (
                    "local_deep_research" in file_path
                    and "site-packages" not in file_path
                    and "venv" not in file_path
                ):
                    # Extract relative path from local_deep_research
                    if "src/local_deep_research" in file_path:
                        relative_path = file_path.split(
                            "src/local_deep_research"
                        )[-1].lstrip("/")
                    elif "local_deep_research/src" in file_path:
                        relative_path = file_path.split(
                            "local_deep_research/src"
                        )[-1].lstrip("/")
                    elif "local_deep_research" in file_path:
                        # Get everything after local_deep_research
                        relative_path = file_path.split("local_deep_research")[
                            -1
                        ].lstrip("/")
                    else:
                        relative_path = Path(file_path).name

                    self.calling_file = relative_path
                    self.calling_function = frame_info.function

                    # Capture a simplified call stack (just the relevant frames)
                    call_stack_frames = []
                    for frame in stack[1:6]:  # Limit to 5 frames
                        if (
                            "local_deep_research" in frame.filename
                            and "site-packages" not in frame.filename
                            and "venv" not in frame.filename
                        ):
                            frame_name = f"{Path(frame.filename).name}:{frame.function}:{frame.lineno}"
                            call_stack_frames.append(frame_name)

                    self.call_stack = (
                        " -> ".join(call_stack_frames)
                        if call_stack_frames
                        else None
                    )
                    break
        except Exception as e:
            logger.debug(f"Error capturing call stack: {e}")
            # Continue without call stack info if there's an error

        # Debug logging
        logger.debug(f"on_llm_start serialized: {serialized}")
        logger.debug(f"on_llm_start kwargs: {kwargs}")

        # First, use preset values if available
        if self.preset_model:
            self.current_model = self.preset_model
        else:
            # Try multiple locations for model name
            model_name = None

            # First check invocation_params
            invocation_params = kwargs.get("invocation_params", {})
            model_name = invocation_params.get(
                "model"
            ) or invocation_params.get("model_name")

            # Check kwargs directly
            if not model_name:
                model_name = kwargs.get("model") or kwargs.get("model_name")

            # Check serialized data
            if not model_name and "kwargs" in serialized:
                model_name = serialized["kwargs"].get("model") or serialized[
                    "kwargs"
                ].get("model_name")

            # Check for name in serialized data
            if not model_name and "name" in serialized:
                model_name = serialized["name"]

            # If still not found and we have Ollama, try to extract from the instance
            if (
                not model_name
                and "_type" in serialized
                and "ChatOllama" in serialized["_type"]
            ):
                # For Ollama, the model name might be in the serialized kwargs
                if "kwargs" in serialized and "model" in serialized["kwargs"]:
                    model_name = serialized["kwargs"]["model"]
                else:
                    # Default to the type if we can't find the actual model
                    model_name = "ollama"

            # Final fallback
            if not model_name:
                if "_type" in serialized:
                    model_name = serialized["_type"]
                else:
                    model_name = "unknown"

            self.current_model = model_name

        # Use preset provider if available
        if self.preset_provider:
            self.current_provider = self.preset_provider
        else:
            # Extract provider from serialized type or kwargs
            if "_type" in serialized:
                type_str = serialized["_type"]
                if "ChatOllama" in type_str:
                    self.current_provider = "ollama"
                elif "ChatOpenAI" in type_str:
                    self.current_provider = "openai"
                elif "ChatAnthropic" in type_str:
                    self.current_provider = "anthropic"
                else:
                    self.current_provider = kwargs.get("provider", "unknown")
            else:
                self.current_provider = kwargs.get("provider", "unknown")

        # Initialize model tracking if needed
        if self.current_model not in self.counts["by_model"]:
            self.counts["by_model"][self.current_model] = {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "calls": 0,
                "provider": self.current_provider,
            }

        # Increment call count
        self.counts["by_model"][self.current_model]["calls"] += 1

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Called when LLM ends running."""
        # Phase 1 Enhancement: Calculate response time
        if self.start_time:
            self.response_time_ms = int((time.time() - self.start_time) * 1000)

        # Extract token usage from response
        token_usage = None

        # Check multiple locations for token usage
        if hasattr(response, "llm_output") and response.llm_output:
            token_usage = response.llm_output.get(
                "token_usage"
            ) or response.llm_output.get("usage", {})

        # Check for usage metadata in generations (Ollama specific)
        if not token_usage and hasattr(response, "generations"):
            for generation_list in response.generations:
                for generation in generation_list:
                    if hasattr(generation, "message") and hasattr(
                        generation.message, "usage_metadata"
                    ):
                        usage_meta = generation.message.usage_metadata
                        if usage_meta:  # Check if usage_metadata is not None
                            token_usage = {
                                "prompt_tokens": usage_meta.get(
                                    "input_tokens", 0
                                ),
                                "completion_tokens": usage_meta.get(
                                    "output_tokens", 0
                                ),
                                "total_tokens": usage_meta.get(
                                    "total_tokens", 0
                                ),
                            }
                            break
                    # Also check response_metadata
                    if hasattr(generation, "message") and hasattr(
                        generation.message, "response_metadata"
                    ):
                        resp_meta = generation.message.response_metadata
                        if resp_meta.get("prompt_eval_count") or resp_meta.get(
                            "eval_count"
                        ):
                            # Capture raw Ollama metrics
                            self.ollama_metrics = {
                                "prompt_eval_count": resp_meta.get(
                                    "prompt_eval_count"
                                ),
                                "eval_count": resp_meta.get("eval_count"),
                                "total_duration": resp_meta.get(
                                    "total_duration"
                                ),
                                "load_duration": resp_meta.get("load_duration"),
                                "prompt_eval_duration": resp_meta.get(
                                    "prompt_eval_duration"
                                ),
                                "eval_duration": resp_meta.get("eval_duration"),
                            }

                            # Check for context overflow
                            prompt_eval_count = resp_meta.get(
                                "prompt_eval_count", 0
                            )
                            if self.context_limit and prompt_eval_count > 0:
                                # Check if we're near or at the context limit
                                if (
                                    prompt_eval_count
                                    >= self.context_limit * 0.95
                                ):  # 95% threshold
                                    self.context_truncated = True

                                    # Estimate tokens truncated
                                    if (
                                        self.original_prompt_estimate
                                        > prompt_eval_count
                                    ):
                                        self.tokens_truncated = max(
                                            0,
                                            self.original_prompt_estimate
                                            - prompt_eval_count,
                                        )
                                        self.truncation_ratio = (
                                            self.tokens_truncated
                                            / self.original_prompt_estimate
                                            if self.original_prompt_estimate > 0
                                            else 0
                                        )
                                        logger.warning(
                                            f"Context overflow detected! "
                                            f"Prompt tokens: {prompt_eval_count}/{self.context_limit} "
                                            f"(estimated {self.tokens_truncated} tokens truncated, "
                                            f"{self.truncation_ratio:.1%} of prompt)"
                                        )

                            token_usage = {
                                "prompt_tokens": resp_meta.get(
                                    "prompt_eval_count", 0
                                ),
                                "completion_tokens": resp_meta.get(
                                    "eval_count", 0
                                ),
                                "total_tokens": resp_meta.get(
                                    "prompt_eval_count", 0
                                )
                                + resp_meta.get("eval_count", 0),
                            }
                            break
                if token_usage:
                    break

        if token_usage and isinstance(token_usage, dict):
            prompt_tokens = token_usage.get("prompt_tokens", 0)
            completion_tokens = token_usage.get("completion_tokens", 0)
            total_tokens = token_usage.get(
                "total_tokens", prompt_tokens + completion_tokens
            )

            # Update in-memory counts
            self.counts["total_prompt_tokens"] += prompt_tokens
            self.counts["total_completion_tokens"] += completion_tokens
            self.counts["total_tokens"] += total_tokens

            if self.current_model:
                self.counts["by_model"][self.current_model][
                    "prompt_tokens"
                ] += prompt_tokens
                self.counts["by_model"][self.current_model][
                    "completion_tokens"
                ] += completion_tokens
                self.counts["by_model"][self.current_model]["total_tokens"] += (
                    total_tokens
                )

            # Save to database if we have a research_id
            if self.research_id:
                self._save_to_db(prompt_tokens, completion_tokens)

    def on_llm_error(self, error, **kwargs: Any) -> None:
        """Called when LLM encounters an error."""
        # Phase 1 Enhancement: Track errors
        if self.start_time:
            self.response_time_ms = int((time.time() - self.start_time) * 1000)

        self.success_status = "error"
        self.error_type = str(type(error).__name__)

        # Still save to database to track failed calls
        if self.research_id:
            self._save_to_db(0, 0)

    def _get_context_overflow_fields(self) -> Dict[str, Any]:
        """Get context overflow detection fields for database saving."""
        return {
            "context_limit": self.context_limit,
            "context_truncated": self.context_truncated,  # Now Boolean
            "tokens_truncated": self.tokens_truncated
            if self.context_truncated
            else None,
            "truncation_ratio": self.truncation_ratio
            if self.context_truncated
            else None,
            # Raw Ollama metrics
            "ollama_prompt_eval_count": self.ollama_metrics.get(
                "prompt_eval_count"
            ),
            "ollama_eval_count": self.ollama_metrics.get("eval_count"),
            "ollama_total_duration": self.ollama_metrics.get("total_duration"),
            "ollama_load_duration": self.ollama_metrics.get("load_duration"),
            "ollama_prompt_eval_duration": self.ollama_metrics.get(
                "prompt_eval_duration"
            ),
            "ollama_eval_duration": self.ollama_metrics.get("eval_duration"),
        }

    def _save_to_db(self, prompt_tokens: int, completion_tokens: int):
        """Save token usage to the database."""
        # Check if we're in a thread - if so, queue the save for later
        import threading

        if threading.current_thread().name != "MainThread":
            # Use thread-safe metrics database for background threads
            username = (
                self.research_context.get("username")
                if self.research_context
                else None
            )

            # Debug logging - use info level to ensure it shows
            logger.info(
                f"Thread metrics save attempt - Thread: {threading.current_thread().name}, "
                f"Username: {username}, Research context keys: {list(self.research_context.keys()) if self.research_context else 'None'}"
            )

            if not username:
                logger.warning(
                    f"Cannot save token metrics - no username in research context. "
                    f"Token usage: prompt={prompt_tokens}, completion={completion_tokens}, "
                    f"Research context: {self.research_context}"
                )
                return

            # Import the thread-safe metrics database

            # Prepare token data
            token_data = {
                "model_name": self.current_model,
                "provider": self.current_provider,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "research_query": self.research_context.get("research_query"),
                "research_mode": self.research_context.get("research_mode"),
                "research_phase": self.research_context.get("research_phase"),
                "search_iteration": self.research_context.get(
                    "search_iteration"
                ),
                "response_time_ms": self.response_time_ms,
                "success_status": self.success_status,
                "error_type": self.error_type,
                "search_engines_planned": self.research_context.get(
                    "search_engines_planned"
                ),
                "search_engine_selected": self.research_context.get(
                    "search_engine_selected"
                ),
                "calling_file": self.calling_file,
                "calling_function": self.calling_function,
                "call_stack": self.call_stack,
                # Add context overflow fields using helper method
                **self._get_context_overflow_fields(),
            }

            # Convert list to JSON string if needed
            if isinstance(token_data.get("search_engines_planned"), list):
                token_data["search_engines_planned"] = json.dumps(
                    token_data["search_engines_planned"]
                )

            # Get password from research context
            password = self.research_context.get("user_password")
            if not password:
                logger.warning(
                    f"Cannot save token metrics - no password in research context. "
                    f"Username: {username}, Token usage: prompt={prompt_tokens}, completion={completion_tokens}"
                )
                return

            # Write metrics directly using thread-safe database
            try:
                from ..database.thread_metrics import metrics_writer

                # Set password for this thread
                metrics_writer.set_user_password(username, password)

                # Write metrics to encrypted database
                metrics_writer.write_token_metrics(
                    username, self.research_id, token_data
                )
                logger.debug(
                    f"Saved token metrics to encrypted DB for user {username} - "
                    f"research {self.research_id} in thread {threading.current_thread().name}. "
                    f"Token usage: prompt={prompt_tokens}, completion={completion_tokens}"
                )
            except Exception:
                logger.exception("Failed to write metrics from thread")
            return

        # In MainThread, save directly
        try:
            from flask import session as flask_session
            from ..database.session_context import get_user_db_session

            username = flask_session.get("username")
            if not username:
                logger.debug("No user session, skipping token metrics save")
                return

            with get_user_db_session(username) as session:
                # Phase 1 Enhancement: Prepare additional context
                research_query = self.research_context.get("research_query")
                research_mode = self.research_context.get("research_mode")
                research_phase = self.research_context.get("research_phase")
                search_iteration = self.research_context.get("search_iteration")
                search_engines_planned = self.research_context.get(
                    "search_engines_planned"
                )
                search_engine_selected = self.research_context.get(
                    "search_engine_selected"
                )

                # Debug logging for search engine context
                if search_engines_planned or search_engine_selected:
                    logger.info(
                        f"Token tracking - Search context: planned={search_engines_planned}, selected={search_engine_selected}, phase={research_phase}"
                    )
                else:
                    logger.debug(
                        f"Token tracking - No search engine context yet, phase={research_phase}"
                    )

                # Convert list to JSON string if needed
                if isinstance(search_engines_planned, list):
                    search_engines_planned = json.dumps(search_engines_planned)

                # Log context overflow detection values before saving
                logger.debug(
                    f"Saving TokenUsage - context_limit: {self.context_limit}, "
                    f"context_truncated: {self.context_truncated}, "
                    f"tokens_truncated: {self.tokens_truncated}, "
                    f"ollama_prompt_eval_count: {self.ollama_metrics.get('prompt_eval_count')}, "
                    f"prompt_tokens: {self.estimated_prompt_tokens}, "
                    f"completion_tokens: {self.estimated_completion_tokens}"
                )

                # Add token usage record with enhanced fields
                token_usage = TokenUsage(
                    research_id=self.research_id,
                    model_name=self.current_model,
                    model_provider=self.current_provider,  # Added provider
                    # for accurate cost tracking
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=prompt_tokens + completion_tokens,
                    # Phase 1 Enhancement: Research context
                    research_query=research_query,
                    research_mode=research_mode,
                    research_phase=research_phase,
                    search_iteration=search_iteration,
                    # Phase 1 Enhancement: Performance metrics
                    response_time_ms=self.response_time_ms,
                    success_status=self.success_status,
                    error_type=self.error_type,
                    # Phase 1 Enhancement: Search engine context
                    search_engines_planned=search_engines_planned,
                    search_engine_selected=search_engine_selected,
                    # Phase 1 Enhancement: Call stack tracking
                    calling_file=self.calling_file,
                    calling_function=self.calling_function,
                    call_stack=self.call_stack,
                    # Add context overflow fields using helper method
                    **self._get_context_overflow_fields(),
                )
                session.add(token_usage)

                # Update or create model usage statistics
                model_usage = (
                    session.query(ModelUsage)
                    .filter_by(
                        model_name=self.current_model,
                    )
                    .first()
                )

                if model_usage:
                    model_usage.total_tokens += (
                        prompt_tokens + completion_tokens
                    )
                    model_usage.total_calls += 1
                else:
                    model_usage = ModelUsage(
                        model_name=self.current_model,
                        model_provider=self.current_provider,
                        total_tokens=prompt_tokens + completion_tokens,
                        total_calls=1,
                    )
                    session.add(model_usage)

                # Commit the transaction
                session.commit()

        except Exception:
            logger.exception("Error saving token usage to database")

    def get_counts(self) -> Dict[str, Any]:
        """Get the current token counts."""
        return self.counts


class TokenCounter:
    """Manager class for token counting across the application."""

    def __init__(self):
        """Initialize the token counter."""
        # No longer need to store database reference
        self._thread_metrics_db = None

    @property
    def thread_metrics_db(self):
        """Lazy load thread metrics writer."""
        if self._thread_metrics_db is None:
            try:
                from ..database.thread_metrics import metrics_writer

                self._thread_metrics_db = metrics_writer
            except ImportError:
                logger.warning("Thread metrics writer not available")
        return self._thread_metrics_db

    def create_callback(
        self,
        research_id: Optional[str] = None,
        research_context: Optional[Dict[str, Any]] = None,
    ) -> TokenCountingCallback:
        """Create a new token counting callback.

        Args:
            research_id: The ID of the research to track tokens for
            research_context: Additional research context for enhanced tracking

        Returns:
            A new TokenCountingCallback instance
        """
        return TokenCountingCallback(
            research_id=research_id, research_context=research_context
        )

    def get_research_metrics(self, research_id: str) -> Dict[str, Any]:
        """Get token metrics for a specific research.

        Args:
            research_id: The ID of the research

        Returns:
            Dictionary containing token usage metrics
        """
        from flask import session as flask_session

        from ..database.session_context import get_user_db_session

        username = flask_session.get("username")
        if not username:
            return {
                "research_id": research_id,
                "total_tokens": 0,
                "total_calls": 0,
                "model_usage": [],
            }

        with get_user_db_session(username) as session:
            # Get token usage for this research from TokenUsage table
            from sqlalchemy import func

            token_usages = (
                session.query(
                    TokenUsage.model_name,
                    TokenUsage.model_provider,
                    func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
                    func.sum(TokenUsage.completion_tokens).label(
                        "completion_tokens"
                    ),
                    func.sum(TokenUsage.total_tokens).label("total_tokens"),
                    func.count().label("calls"),
                )
                .filter_by(research_id=research_id)
                .group_by(TokenUsage.model_name, TokenUsage.model_provider)
                .order_by(func.sum(TokenUsage.total_tokens).desc())
                .all()
            )

            model_usage = []
            total_tokens = 0
            total_calls = 0

            for usage in token_usages:
                model_usage.append(
                    {
                        "model": usage.model_name,
                        "provider": usage.model_provider,
                        "tokens": usage.total_tokens or 0,
                        "calls": usage.calls or 0,
                        "prompt_tokens": usage.prompt_tokens or 0,
                        "completion_tokens": usage.completion_tokens or 0,
                    }
                )
                total_tokens += usage.total_tokens or 0
                total_calls += usage.calls or 0

            return {
                "research_id": research_id,
                "total_tokens": total_tokens,
                "total_calls": total_calls,
                "model_usage": model_usage,
            }

    def get_overall_metrics(
        self, period: str = "30d", research_mode: str = "all"
    ) -> Dict[str, Any]:
        """Get overall token metrics across all researches.

        Args:
            period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
            research_mode: Research mode to filter by ('quick', 'detailed', 'all')

        Returns:
            Dictionary containing overall metrics
        """
        # First get metrics from user's encrypted database
        encrypted_metrics = self._get_metrics_from_encrypted_db(
            period, research_mode
        )

        # Then get metrics from thread-safe metrics database
        thread_metrics = self._get_metrics_from_thread_db(period, research_mode)

        # Merge the results
        return self._merge_metrics(encrypted_metrics, thread_metrics)

    def _get_metrics_from_encrypted_db(
        self, period: str, research_mode: str
    ) -> Dict[str, Any]:
        """Get metrics from user's encrypted database."""
        from flask import session as flask_session

        from ..database.session_context import get_user_db_session

        username = flask_session.get("username")
        if not username:
            return self._get_empty_metrics()

        try:
            with get_user_db_session(username) as session:
                # Build base query with filters
                query = session.query(TokenUsage)

                # Apply time filter
                time_condition = get_time_filter_condition(
                    period, TokenUsage.timestamp
                )
                if time_condition is not None:
                    query = query.filter(time_condition)

                # Apply research mode filter
                mode_condition = get_research_mode_condition(
                    research_mode, TokenUsage.research_mode
                )
                if mode_condition is not None:
                    query = query.filter(mode_condition)

                # Total tokens from TokenUsage
                total_tokens = (
                    query.with_entities(
                        func.sum(TokenUsage.total_tokens)
                    ).scalar()
                    or 0
                )

                # Import ResearchHistory model
                from ..database.models.research import ResearchHistory

                # Count researches from ResearchHistory table
                research_query = session.query(func.count(ResearchHistory.id))

                # Debug: Check if any research history records exist at all
                all_research_count = (
                    session.query(func.count(ResearchHistory.id)).scalar() or 0
                )
                logger.warning(
                    f"DEBUG: Total ResearchHistory records in database: {all_research_count}"
                )

                # Debug: List first few research IDs and their timestamps
                sample_researches = (
                    session.query(
                        ResearchHistory.id,
                        ResearchHistory.created_at,
                        ResearchHistory.mode,
                    )
                    .limit(5)
                    .all()
                )
                if sample_researches:
                    logger.warning("DEBUG: Sample ResearchHistory records:")
                    for r_id, r_created, r_mode in sample_researches:
                        logger.warning(
                            f"  - ID: {r_id}, Created: {r_created}, Mode: {r_mode}"
                        )
                else:
                    logger.warning(
                        "DEBUG: No ResearchHistory records found in database!"
                    )

                # Get time filter conditions for ResearchHistory query
                start_time, end_time = None, None
                if period != "all":
                    if period == "today":
                        start_time = datetime.now(UTC).replace(
                            hour=0, minute=0, second=0, microsecond=0
                        )
                    elif period == "week":
                        start_time = datetime.now(UTC) - timedelta(days=7)
                    elif period == "month":
                        start_time = datetime.now(UTC) - timedelta(days=30)

                    if start_time:
                        end_time = datetime.now(UTC)

                # Apply time filter if specified
                if start_time and end_time:
                    research_query = research_query.filter(
                        ResearchHistory.created_at >= start_time.isoformat(),
                        ResearchHistory.created_at <= end_time.isoformat(),
                    )

                # Apply mode filter if specified
                mode_filter = research_mode if research_mode != "all" else None
                if mode_filter:
                    logger.warning(
                        f"DEBUG: Applying mode filter: {mode_filter}"
                    )
                    research_query = research_query.filter(
                        ResearchHistory.mode == mode_filter
                    )

                total_researches = research_query.scalar() or 0
                logger.warning(
                    f"DEBUG: Final filtered research count: {total_researches}"
                )

                # Also check distinct research_ids in TokenUsage for comparison
                token_research_count = (
                    session.query(
                        func.count(func.distinct(TokenUsage.research_id))
                    ).scalar()
                    or 0
                )
                logger.warning(
                    f"DEBUG: Distinct research_ids in TokenUsage: {token_research_count}"
                )

                # Model statistics using ORM aggregation
                model_stats_query = session.query(
                    TokenUsage.model_name,
                    func.sum(TokenUsage.total_tokens).label("tokens"),
                    func.count().label("calls"),
                    func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
                    func.sum(TokenUsage.completion_tokens).label(
                        "completion_tokens"
                    ),
                ).filter(TokenUsage.model_name.isnot(None))

                # Apply same filters to model stats
                if time_condition is not None:
                    model_stats_query = model_stats_query.filter(time_condition)
                if mode_condition is not None:
                    model_stats_query = model_stats_query.filter(mode_condition)

                model_stats = (
                    model_stats_query.group_by(TokenUsage.model_name)
                    .order_by(func.sum(TokenUsage.total_tokens).desc())
                    .all()
                )

                # Get provider info from ModelUsage table
                by_model = []
                for stat in model_stats:
                    # Try to get provider from ModelUsage table
                    provider_info = (
                        session.query(ModelUsage.model_provider)
                        .filter(ModelUsage.model_name == stat.model_name)
                        .first()
                    )
                    provider = (
                        provider_info.model_provider
                        if provider_info
                        else "unknown"
                    )

                    by_model.append(
                        {
                            "model": stat.model_name,
                            "provider": provider,
                            "tokens": stat.tokens,
                            "calls": stat.calls,
                            "prompt_tokens": stat.prompt_tokens,
                            "completion_tokens": stat.completion_tokens,
                        }
                    )

                # Get recent researches with token usage
                # Note: This requires research_history table - for now we'll use available data
                recent_research_query = session.query(
                    TokenUsage.research_id,
                    func.sum(TokenUsage.total_tokens).label("token_count"),
                    func.max(TokenUsage.timestamp).label("latest_timestamp"),
                ).filter(TokenUsage.research_id.isnot(None))

                if time_condition is not None:
                    recent_research_query = recent_research_query.filter(
                        time_condition
                    )
                if mode_condition is not None:
                    recent_research_query = recent_research_query.filter(
                        mode_condition
                    )

                recent_research_data = (
                    recent_research_query.group_by(TokenUsage.research_id)
                    .order_by(func.max(TokenUsage.timestamp).desc())
                    .limit(10)
                    .all()
                )

                recent_researches = []
                for research_data in recent_research_data:
                    # Get research query from token_usage table if available
                    research_query_data = (
                        session.query(TokenUsage.research_query)
                        .filter(
                            TokenUsage.research_id == research_data.research_id,
                            TokenUsage.research_query.isnot(None),
                        )
                        .first()
                    )

                    query_text = (
                        research_query_data.research_query
                        if research_query_data
                        else f"Research {research_data.research_id}"
                    )

                    recent_researches.append(
                        {
                            "id": research_data.research_id,
                            "query": query_text,
                            "tokens": research_data.token_count or 0,
                            "created_at": research_data.latest_timestamp,
                        }
                    )

                # Token breakdown statistics
                breakdown_query = query.with_entities(
                    func.sum(TokenUsage.prompt_tokens).label(
                        "total_input_tokens"
                    ),
                    func.sum(TokenUsage.completion_tokens).label(
                        "total_output_tokens"
                    ),
                    func.avg(TokenUsage.prompt_tokens).label(
                        "avg_input_tokens"
                    ),
                    func.avg(TokenUsage.completion_tokens).label(
                        "avg_output_tokens"
                    ),
                    func.avg(TokenUsage.total_tokens).label("avg_total_tokens"),
                )
                token_breakdown = breakdown_query.first()

                # Get rate limiting metrics
                from ..database.models import (
                    RateLimitAttempt,
                    RateLimitEstimate,
                )

                # Get rate limit attempts
                rate_limit_query = session.query(RateLimitAttempt)

                # Apply time filter
                if time_condition is not None:
                    # RateLimitAttempt uses timestamp as float, not datetime
                    if period == "7d":
                        cutoff_time = time.time() - (7 * 24 * 3600)
                    elif period == "30d":
                        cutoff_time = time.time() - (30 * 24 * 3600)
                    elif period == "3m":
                        cutoff_time = time.time() - (90 * 24 * 3600)
                    elif period == "1y":
                        cutoff_time = time.time() - (365 * 24 * 3600)
                    else:  # all
                        cutoff_time = 0

                    if cutoff_time > 0:
                        rate_limit_query = rate_limit_query.filter(
                            RateLimitAttempt.timestamp >= cutoff_time
                        )

                # Get rate limit statistics
                total_attempts = rate_limit_query.count()
                successful_attempts = rate_limit_query.filter(
                    RateLimitAttempt.success
                ).count()
                failed_attempts = total_attempts - successful_attempts

                # Count rate limiting events (failures with RateLimitError)
                rate_limit_events = rate_limit_query.filter(
                    ~RateLimitAttempt.success,
                    RateLimitAttempt.error_type == "RateLimitError",
                ).count()

                logger.warning(
                    f"DEBUG: Rate limit attempts in database: total={total_attempts}, successful={successful_attempts}"
                )

                # Get all attempts for detailed calculations
                attempts = rate_limit_query.all()

                # Calculate average wait times
                if attempts:
                    avg_wait_time = sum(a.wait_time for a in attempts) / len(
                        attempts
                    )
                    successful_wait_times = [
                        a.wait_time for a in attempts if a.success
                    ]
                    avg_successful_wait = (
                        sum(successful_wait_times) / len(successful_wait_times)
                        if successful_wait_times
                        else 0
                    )
                else:
                    avg_wait_time = 0
                    avg_successful_wait = 0

                # Get tracked engines - count distinct engine types from attempts
                tracked_engines_query = session.query(
                    func.count(func.distinct(RateLimitAttempt.engine_type))
                )
                if cutoff_time > 0:
                    tracked_engines_query = tracked_engines_query.filter(
                        RateLimitAttempt.timestamp >= cutoff_time
                    )
                tracked_engines = tracked_engines_query.scalar() or 0

                # Get engine-specific stats from attempts
                engine_stats = []

                # Get distinct engine types from attempts
                engine_types_query = session.query(
                    RateLimitAttempt.engine_type
                ).distinct()
                if cutoff_time > 0:
                    engine_types_query = engine_types_query.filter(
                        RateLimitAttempt.timestamp >= cutoff_time
                    )
                engine_types = [
                    row.engine_type for row in engine_types_query.all()
                ]

                for engine_type in engine_types:
                    engine_attempts_list = [
                        a for a in attempts if a.engine_type == engine_type
                    ]
                    engine_attempts = len(engine_attempts_list)
                    engine_success = len(
                        [a for a in engine_attempts_list if a.success]
                    )

                    # Get estimate if exists
                    estimate = (
                        session.query(RateLimitEstimate)
                        .filter(RateLimitEstimate.engine_type == engine_type)
                        .first()
                    )

                    # Calculate recent success rate
                    recent_success_rate = (
                        (engine_success / engine_attempts * 100)
                        if engine_attempts > 0
                        else 0
                    )

                    # Determine status based on success rate
                    if estimate:
                        status = (
                            "healthy"
                            if estimate.success_rate > 0.8
                            else "degraded"
                            if estimate.success_rate > 0.5
                            else "poor"
                        )
                    else:
                        status = (
                            "healthy"
                            if recent_success_rate > 80
                            else "degraded"
                            if recent_success_rate > 50
                            else "poor"
                        )

                    engine_stat = {
                        "engine": engine_type,
                        "base_wait": estimate.base_wait_seconds
                        if estimate
                        else 0.0,
                        "base_wait_seconds": round(
                            estimate.base_wait_seconds if estimate else 0.0, 2
                        ),
                        "min_wait_seconds": round(
                            estimate.min_wait_seconds if estimate else 0.0, 2
                        ),
                        "max_wait_seconds": round(
                            estimate.max_wait_seconds if estimate else 0.0, 2
                        ),
                        "success_rate": round(estimate.success_rate * 100, 1)
                        if estimate
                        else recent_success_rate,
                        "total_attempts": estimate.total_attempts
                        if estimate
                        else engine_attempts,
                        "recent_attempts": engine_attempts,
                        "recent_success_rate": round(recent_success_rate, 1),
                        "attempts": engine_attempts,
                        "status": status,
                    }

                    if estimate:
                        engine_stat["last_updated"] = datetime.fromtimestamp(
                            estimate.last_updated
                        ).strftime("%Y-%m-%d %H:%M:%S")
                    else:
                        engine_stat["last_updated"] = "Never"

                    engine_stats.append(engine_stat)

                logger.warning(
                    f"DEBUG: Tracked engines: {tracked_engines}, engine_stats: {engine_stats}"
                )

                result = {
                    "total_tokens": total_tokens,
                    "total_researches": total_researches,
                    "by_model": by_model,
                    "recent_researches": recent_researches,
                    "token_breakdown": {
                        "total_input_tokens": int(
                            token_breakdown.total_input_tokens or 0
                        ),
                        "total_output_tokens": int(
                            token_breakdown.total_output_tokens or 0
                        ),
                        "avg_input_tokens": int(
                            token_breakdown.avg_input_tokens or 0
                        ),
                        "avg_output_tokens": int(
                            token_breakdown.avg_output_tokens or 0
                        ),
                        "avg_total_tokens": int(
                            token_breakdown.avg_total_tokens or 0
                        ),
                    },
                    "rate_limiting": {
                        "total_attempts": total_attempts,
                        "successful_attempts": successful_attempts,
                        "failed_attempts": failed_attempts,
                        "success_rate": (
                            successful_attempts / total_attempts * 100
                        )
                        if total_attempts > 0
                        else 0,
                        "rate_limit_events": rate_limit_events,
                        "avg_wait_time": round(float(avg_wait_time), 2),
                        "avg_successful_wait": round(
                            float(avg_successful_wait), 2
                        ),
                        "tracked_engines": tracked_engines,
                        "engine_stats": engine_stats,
                        "total_engines_tracked": tracked_engines,
                        "healthy_engines": len(
                            [
                                s
                                for s in engine_stats
                                if s["status"] == "healthy"
                            ]
                        ),
                        "degraded_engines": len(
                            [
                                s
                                for s in engine_stats
                                if s["status"] == "degraded"
                            ]
                        ),
                        "poor_engines": len(
                            [s for s in engine_stats if s["status"] == "poor"]
                        ),
                    },
                }

                logger.warning(
                    f"DEBUG: Returning from _get_metrics_from_encrypted_db - total_researches: {result['total_researches']}"
                )
                return result
        except Exception as e:
            logger.exception(
                f"CRITICAL ERROR accessing encrypted database for metrics: {e}"
            )
            return self._get_empty_metrics()

    def _get_metrics_from_thread_db(
        self, period: str, research_mode: str
    ) -> Dict[str, Any]:
        """Get metrics from thread-safe metrics database."""
        if not self.thread_metrics_db:
            return {
                "total_tokens": 0,
                "total_researches": 0,
                "by_model": [],
                "recent_researches": [],
                "token_breakdown": {
                    "total_input_tokens": 0,
                    "total_output_tokens": 0,
                    "avg_input_tokens": 0,
                    "avg_output_tokens": 0,
                    "avg_total_tokens": 0,
                },
            }

        try:
            with self.thread_metrics_db.get_session() as session:
                # Build base query with filters
                query = session.query(TokenUsage)

                # Apply time filter
                time_condition = get_time_filter_condition(
                    period, TokenUsage.timestamp
                )
                if time_condition is not None:
                    query = query.filter(time_condition)

                # Apply research mode filter
                mode_condition = get_research_mode_condition(
                    research_mode, TokenUsage.research_mode
                )
                if mode_condition is not None:
                    query = query.filter(mode_condition)

                # Get totals
                total_tokens = (
                    query.with_entities(
                        func.sum(TokenUsage.total_tokens)
                    ).scalar()
                    or 0
                )
                total_researches = (
                    query.with_entities(
                        func.count(func.distinct(TokenUsage.research_id))
                    ).scalar()
                    or 0
                )

                # Get model statistics
                model_stats = (
                    query.with_entities(
                        TokenUsage.model_name,
                        func.sum(TokenUsage.total_tokens).label("tokens"),
                        func.count().label("calls"),
                        func.sum(TokenUsage.prompt_tokens).label(
                            "prompt_tokens"
                        ),
                        func.sum(TokenUsage.completion_tokens).label(
                            "completion_tokens"
                        ),
                    )
                    .filter(TokenUsage.model_name.isnot(None))
                    .group_by(TokenUsage.model_name)
                    .all()
                )

                by_model = []
                for stat in model_stats:
                    by_model.append(
                        {
                            "model": stat.model_name,
                            "provider": "unknown",  # Provider info might not be in thread DB
                            "tokens": stat.tokens,
                            "calls": stat.calls,
                            "prompt_tokens": stat.prompt_tokens,
                            "completion_tokens": stat.completion_tokens,
                        }
                    )

                # Token breakdown
                breakdown = query.with_entities(
                    func.sum(TokenUsage.prompt_tokens).label(
                        "total_input_tokens"
                    ),
                    func.sum(TokenUsage.completion_tokens).label(
                        "total_output_tokens"
                    ),
                ).first()

                return {
                    "total_tokens": total_tokens,
                    "total_researches": total_researches,
                    "by_model": by_model,
                    "recent_researches": [],  # Skip for thread DB
                    "token_breakdown": {
                        "total_input_tokens": int(
                            breakdown.total_input_tokens or 0
                        ),
                        "total_output_tokens": int(
                            breakdown.total_output_tokens or 0
                        ),
                        "avg_input_tokens": 0,
                        "avg_output_tokens": 0,
                        "avg_total_tokens": 0,
                    },
                }
        except Exception:
            logger.exception("Error reading thread metrics database")
            return {
                "total_tokens": 0,
                "total_researches": 0,
                "by_model": [],
                "recent_researches": [],
                "token_breakdown": {
                    "total_input_tokens": 0,
                    "total_output_tokens": 0,
                    "avg_input_tokens": 0,
                    "avg_output_tokens": 0,
                    "avg_total_tokens": 0,
                },
            }

    def _merge_metrics(
        self, encrypted: Dict[str, Any], thread: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Merge metrics from both databases."""
        # Combine totals
        total_tokens = encrypted.get("total_tokens", 0) + thread.get(
            "total_tokens", 0
        )
        total_researches = max(
            encrypted.get("total_researches", 0),
            thread.get("total_researches", 0),
        )
        logger.warning(
            f"DEBUG: Merged metrics - encrypted researches: {encrypted.get('total_researches', 0)}, thread researches: {thread.get('total_researches', 0)}, final: {total_researches}"
        )

        # Merge model usage
        model_map = {}
        for model_data in encrypted.get("by_model", []):
            key = model_data["model"]
            model_map[key] = model_data

        for model_data in thread.get("by_model", []):
            key = model_data["model"]
            if key in model_map:
                # Merge with existing
                model_map[key]["tokens"] += model_data["tokens"]
                model_map[key]["calls"] += model_data["calls"]
                model_map[key]["prompt_tokens"] += model_data["prompt_tokens"]
                model_map[key]["completion_tokens"] += model_data[
                    "completion_tokens"
                ]
            else:
                model_map[key] = model_data

        by_model = sorted(
            model_map.values(), key=lambda x: x["tokens"], reverse=True
        )

        # Merge token breakdown
        token_breakdown = {
            "total_input_tokens": (
                encrypted.get("token_breakdown", {}).get(
                    "total_input_tokens", 0
                )
                + thread.get("token_breakdown", {}).get("total_input_tokens", 0)
            ),
            "total_output_tokens": (
                encrypted.get("token_breakdown", {}).get(
                    "total_output_tokens", 0
                )
                + thread.get("token_breakdown", {}).get(
                    "total_output_tokens", 0
                )
            ),
            "avg_input_tokens": encrypted.get("token_breakdown", {}).get(
                "avg_input_tokens", 0
            ),
            "avg_output_tokens": encrypted.get("token_breakdown", {}).get(
                "avg_output_tokens", 0
            ),
            "avg_total_tokens": encrypted.get("token_breakdown", {}).get(
                "avg_total_tokens", 0
            ),
        }

        result = {
            "total_tokens": total_tokens,
            "total_researches": total_researches,
            "by_model": by_model,
            "recent_researches": encrypted.get("recent_researches", []),
            "token_breakdown": token_breakdown,
        }

        logger.warning(
            f"DEBUG: Final get_token_metrics result - total_researches: {result['total_researches']}"
        )
        return result

    def _get_empty_metrics(self) -> Dict[str, Any]:
        """Return empty metrics structure when no data is available."""
        return {
            "total_tokens": 0,
            "total_researches": 0,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "avg_prompt_tokens": 0,
                "avg_completion_tokens": 0,
                "avg_total_tokens": 0,
            },
        }

    def get_enhanced_metrics(
        self, period: str = "30d", research_mode: str = "all"
    ) -> Dict[str, Any]:
        """Get enhanced Phase 1 tracking metrics.

        Args:
            period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
            research_mode: Research mode to filter by ('quick', 'detailed', 'all')

        Returns:
            Dictionary containing enhanced metrics data including time series
        """
        from flask import session as flask_session

        from ..database.session_context import get_user_db_session

        username = flask_session.get("username")
        if not username:
            # Return empty metrics structure when no user session
            return {
                "recent_enhanced_data": [],
                "performance_stats": {
                    "avg_response_time": 0,
                    "min_response_time": 0,
                    "max_response_time": 0,
                    "success_rate": 0,
                    "error_rate": 0,
                    "total_enhanced_calls": 0,
                },
                "mode_breakdown": [],
                "search_engine_stats": [],
                "phase_breakdown": [],
                "time_series_data": [],
                "call_stack_analysis": {
                    "by_file": [],
                    "by_function": [],
                },
            }

        try:
            with get_user_db_session(username) as session:
                # Build base query with filters
                query = session.query(TokenUsage)

                # Apply time filter
                time_condition = get_time_filter_condition(
                    period, TokenUsage.timestamp
                )
                if time_condition is not None:
                    query = query.filter(time_condition)

                # Apply research mode filter
                mode_condition = get_research_mode_condition(
                    research_mode, TokenUsage.research_mode
                )
                if mode_condition is not None:
                    query = query.filter(mode_condition)

                # Get time series data for the chart - most important for "Token Consumption Over Time"
                time_series_query = query.filter(
                    TokenUsage.timestamp.isnot(None),
                    TokenUsage.total_tokens > 0,
                ).order_by(TokenUsage.timestamp.asc())

                # Limit to recent data for performance
                if period != "all":
                    time_series_query = time_series_query.limit(200)

                time_series_data = time_series_query.all()

                # Format time series data with cumulative calculations
                time_series = []
                cumulative_tokens = 0
                cumulative_prompt_tokens = 0
                cumulative_completion_tokens = 0

                for usage in time_series_data:
                    cumulative_tokens += usage.total_tokens or 0
                    cumulative_prompt_tokens += usage.prompt_tokens or 0
                    cumulative_completion_tokens += usage.completion_tokens or 0

                    time_series.append(
                        {
                            "timestamp": str(usage.timestamp)
                            if usage.timestamp
                            else None,
                            "tokens": usage.total_tokens or 0,
                            "prompt_tokens": usage.prompt_tokens or 0,
                            "completion_tokens": usage.completion_tokens or 0,
                            "cumulative_tokens": cumulative_tokens,
                            "cumulative_prompt_tokens": cumulative_prompt_tokens,
                            "cumulative_completion_tokens": cumulative_completion_tokens,
                            "research_id": usage.research_id,
                        }
                    )

                # Basic performance stats using ORM
                performance_query = query.filter(
                    TokenUsage.response_time_ms.isnot(None)
                )
                total_calls = performance_query.count()

                if total_calls > 0:
                    avg_response_time = (
                        performance_query.with_entities(
                            func.avg(TokenUsage.response_time_ms)
                        ).scalar()
                        or 0
                    )
                    min_response_time = (
                        performance_query.with_entities(
                            func.min(TokenUsage.response_time_ms)
                        ).scalar()
                        or 0
                    )
                    max_response_time = (
                        performance_query.with_entities(
                            func.max(TokenUsage.response_time_ms)
                        ).scalar()
                        or 0
                    )
                    success_count = performance_query.filter(
                        TokenUsage.success_status == "success"
                    ).count()
                    error_count = performance_query.filter(
                        TokenUsage.success_status == "error"
                    ).count()

                    perf_stats = {
                        "avg_response_time": round(avg_response_time),
                        "min_response_time": min_response_time,
                        "max_response_time": max_response_time,
                        "success_rate": (
                            round((success_count / total_calls * 100), 1)
                            if total_calls > 0
                            else 0
                        ),
                        "error_rate": (
                            round((error_count / total_calls * 100), 1)
                            if total_calls > 0
                            else 0
                        ),
                        "total_enhanced_calls": total_calls,
                    }
                else:
                    perf_stats = {
                        "avg_response_time": 0,
                        "min_response_time": 0,
                        "max_response_time": 0,
                        "success_rate": 0,
                        "error_rate": 0,
                        "total_enhanced_calls": 0,
                    }

                # Research mode breakdown using ORM
                mode_stats = (
                    query.filter(TokenUsage.research_mode.isnot(None))
                    .with_entities(
                        TokenUsage.research_mode,
                        func.count().label("count"),
                        func.avg(TokenUsage.total_tokens).label("avg_tokens"),
                        func.avg(TokenUsage.response_time_ms).label(
                            "avg_response_time"
                        ),
                    )
                    .group_by(TokenUsage.research_mode)
                    .all()
                )

                modes = [
                    {
                        "mode": stat.research_mode,
                        "count": stat.count,
                        "avg_tokens": round(stat.avg_tokens or 0),
                        "avg_response_time": round(stat.avg_response_time or 0),
                    }
                    for stat in mode_stats
                ]

                # Recent enhanced data (simplified)
                recent_enhanced_query = (
                    query.filter(TokenUsage.research_query.isnot(None))
                    .order_by(TokenUsage.timestamp.desc())
                    .limit(50)
                )

                recent_enhanced_data = recent_enhanced_query.all()
                recent_enhanced = [
                    {
                        "research_query": usage.research_query,
                        "research_mode": usage.research_mode,
                        "research_phase": usage.research_phase,
                        "search_iteration": usage.search_iteration,
                        "response_time_ms": usage.response_time_ms,
                        "success_status": usage.success_status,
                        "error_type": usage.error_type,
                        "search_engines_planned": usage.search_engines_planned,
                        "search_engine_selected": usage.search_engine_selected,
                        "total_tokens": usage.total_tokens,
                        "prompt_tokens": usage.prompt_tokens,
                        "completion_tokens": usage.completion_tokens,
                        "timestamp": str(usage.timestamp)
                        if usage.timestamp
                        else None,
                        "research_id": usage.research_id,
                        "calling_file": usage.calling_file,
                        "calling_function": usage.calling_function,
                        "call_stack": usage.call_stack,
                    }
                    for usage in recent_enhanced_data
                ]

                # Search engine breakdown using ORM
                search_engine_stats = (
                    query.filter(TokenUsage.search_engine_selected.isnot(None))
                    .with_entities(
                        TokenUsage.search_engine_selected,
                        func.count().label("count"),
                        func.avg(TokenUsage.total_tokens).label("avg_tokens"),
                        func.avg(TokenUsage.response_time_ms).label(
                            "avg_response_time"
                        ),
                    )
                    .group_by(TokenUsage.search_engine_selected)
                    .all()
                )

                search_engines = [
                    {
                        "search_engine": stat.search_engine_selected,
                        "count": stat.count,
                        "avg_tokens": round(stat.avg_tokens or 0),
                        "avg_response_time": round(stat.avg_response_time or 0),
                    }
                    for stat in search_engine_stats
                ]

                # Research phase breakdown using ORM
                phase_stats = (
                    query.filter(TokenUsage.research_phase.isnot(None))
                    .with_entities(
                        TokenUsage.research_phase,
                        func.count().label("count"),
                        func.avg(TokenUsage.total_tokens).label("avg_tokens"),
                        func.avg(TokenUsage.response_time_ms).label(
                            "avg_response_time"
                        ),
                    )
                    .group_by(TokenUsage.research_phase)
                    .all()
                )

                phases = [
                    {
                        "phase": stat.research_phase,
                        "count": stat.count,
                        "avg_tokens": round(stat.avg_tokens or 0),
                        "avg_response_time": round(stat.avg_response_time or 0),
                    }
                    for stat in phase_stats
                ]

                # Call stack analysis using ORM
                file_stats = (
                    query.filter(TokenUsage.calling_file.isnot(None))
                    .with_entities(
                        TokenUsage.calling_file,
                        func.count().label("count"),
                        func.avg(TokenUsage.total_tokens).label("avg_tokens"),
                    )
                    .group_by(TokenUsage.calling_file)
                    .order_by(func.count().desc())
                    .limit(10)
                    .all()
                )

                files = [
                    {
                        "file": stat.calling_file,
                        "count": stat.count,
                        "avg_tokens": round(stat.avg_tokens or 0),
                    }
                    for stat in file_stats
                ]

                function_stats = (
                    query.filter(TokenUsage.calling_function.isnot(None))
                    .with_entities(
                        TokenUsage.calling_function,
                        func.count().label("count"),
                        func.avg(TokenUsage.total_tokens).label("avg_tokens"),
                    )
                    .group_by(TokenUsage.calling_function)
                    .order_by(func.count().desc())
                    .limit(10)
                    .all()
                )

                functions = [
                    {
                        "function": stat.calling_function,
                        "count": stat.count,
                        "avg_tokens": round(stat.avg_tokens or 0),
                    }
                    for stat in function_stats
                ]

                return {
                    "recent_enhanced_data": recent_enhanced,
                    "performance_stats": perf_stats,
                    "mode_breakdown": modes,
                    "search_engine_stats": search_engines,
                    "phase_breakdown": phases,
                    "time_series_data": time_series,
                    "call_stack_analysis": {
                        "by_file": files,
                        "by_function": functions,
                    },
                }
        except Exception:
            logger.exception("Error in get_enhanced_metrics")
            # Return simplified response without non-existent columns
            return {
                "recent_enhanced_data": [],
                "performance_stats": {
                    "avg_response_time": 0,
                    "min_response_time": 0,
                    "max_response_time": 0,
                    "success_rate": 0,
                    "error_rate": 0,
                    "total_enhanced_calls": 0,
                },
                "mode_breakdown": [],
                "search_engine_stats": [],
                "phase_breakdown": [],
                "time_series_data": [],
                "call_stack_analysis": {
                    "by_file": [],
                    "by_function": [],
                },
            }

    def get_research_timeline_metrics(self, research_id: str) -> Dict[str, Any]:
        """Get timeline metrics for a specific research.

        Args:
            research_id: The ID of the research

        Returns:
            Dictionary containing timeline metrics for the research
        """
        from flask import session as flask_session

        from ..database.session_context import get_user_db_session

        username = flask_session.get("username")
        if not username:
            return {
                "research_id": research_id,
                "research_details": {},
                "timeline": [],
                "summary": {
                    "total_calls": 0,
                    "total_tokens": 0,
                    "total_prompt_tokens": 0,
                    "total_completion_tokens": 0,
                    "avg_response_time": 0,
                    "success_rate": 0,
                },
                "phase_stats": {},
            }

        with get_user_db_session(username) as session:
            # Get all token usage for this research ordered by time including call stack
            timeline_data = session.execute(
                text(
                    """
                SELECT
                    timestamp,
                    total_tokens,
                    prompt_tokens,
                    completion_tokens,
                    response_time_ms,
                    success_status,
                    error_type,
                    research_phase,
                    search_iteration,
                    search_engine_selected,
                    model_name,
                    calling_file,
                    calling_function,
                    call_stack
                FROM token_usage
                WHERE research_id = :research_id
                ORDER BY timestamp ASC
            """
                ),
                {"research_id": research_id},
            ).fetchall()

            # Format timeline data with cumulative tokens
            timeline = []
            cumulative_tokens = 0
            cumulative_prompt_tokens = 0
            cumulative_completion_tokens = 0

            for row in timeline_data:
                cumulative_tokens += row[1] or 0
                cumulative_prompt_tokens += row[2] or 0
                cumulative_completion_tokens += row[3] or 0

                timeline.append(
                    {
                        "timestamp": str(row[0]) if row[0] else None,
                        "tokens": row[1] or 0,
                        "prompt_tokens": row[2] or 0,
                        "completion_tokens": row[3] or 0,
                        "cumulative_tokens": cumulative_tokens,
                        "cumulative_prompt_tokens": cumulative_prompt_tokens,
                        "cumulative_completion_tokens": cumulative_completion_tokens,
                        "response_time_ms": row[4],
                        "success_status": row[5],
                        "error_type": row[6],
                        "research_phase": row[7],
                        "search_iteration": row[8],
                        "search_engine_selected": row[9],
                        "model_name": row[10],
                        "calling_file": row[11],
                        "calling_function": row[12],
                        "call_stack": row[13],
                    }
                )

            # Get research basic info
            research_info = session.execute(
                text(
                    """
                SELECT query, mode, status, created_at, completed_at
                FROM research_history
                WHERE id = :research_id
            """
                ),
                {"research_id": research_id},
            ).fetchone()

            research_details = {}
            if research_info:
                research_details = {
                    "query": research_info[0],
                    "mode": research_info[1],
                    "status": research_info[2],
                    "created_at": str(research_info[3])
                    if research_info[3]
                    else None,
                    "completed_at": str(research_info[4])
                    if research_info[4]
                    else None,
                }

            # Calculate summary stats
            total_calls = len(timeline_data)
            total_tokens = cumulative_tokens
            avg_response_time = sum(row[4] or 0 for row in timeline_data) / max(
                total_calls, 1
            )
            success_rate = (
                sum(1 for row in timeline_data if row[5] == "success")
                / max(total_calls, 1)
                * 100
            )

            # Phase breakdown for this research
            phase_stats = {}
            for row in timeline_data:
                phase = row[7] or "unknown"
                if phase not in phase_stats:
                    phase_stats[phase] = {
                        "count": 0,
                        "tokens": 0,
                        "avg_response_time": 0,
                    }
                phase_stats[phase]["count"] += 1
                phase_stats[phase]["tokens"] += row[1] or 0
                if row[4]:
                    phase_stats[phase]["avg_response_time"] += row[4]

            # Calculate averages for phases
            for phase in phase_stats:
                if phase_stats[phase]["count"] > 0:
                    phase_stats[phase]["avg_response_time"] = round(
                        phase_stats[phase]["avg_response_time"]
                        / phase_stats[phase]["count"]
                    )

            return {
                "research_id": research_id,
                "research_details": research_details,
                "timeline": timeline,
                "summary": {
                    "total_calls": total_calls,
                    "total_tokens": total_tokens,
                    "total_prompt_tokens": cumulative_prompt_tokens,
                    "total_completion_tokens": cumulative_completion_tokens,
                    "avg_response_time": round(avg_response_time),
                    "success_rate": round(success_rate, 1),
                },
                "phase_stats": phase_stats,
            }
