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

import inspect
import json
import os
import time
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 import get_metrics_db
from .db_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[int] = 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

        # 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()

        # 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 = os.path.basename(file_path)

                    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"{os.path.basename(frame.filename)}:{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
                        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
                    elif 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"
                        ):
                            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:
            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 _save_to_db(self, prompt_tokens: int, completion_tokens: int):
        """Save token usage to the database."""
        try:
            db = get_metrics_db()
            with db.get_session() 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)

                # Add token usage record with enhanced fields
                token_usage = TokenUsage(
                    research_id=self.research_id,
                    model_name=self.current_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,
                )
                session.add(token_usage)

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

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

        except Exception as e:
            logger.exception(f"Error saving token usage to database: {e}")

    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."""
        self.db = get_metrics_db()

    def create_callback(
        self,
        research_id: Optional[int] = 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: int) -> Dict[str, Any]:
        """Get token metrics for a specific research.

        Args:
            research_id: The ID of the research

        Returns:
            Dictionary containing token usage metrics
        """
        with self.db.get_session() as session:
            # Get model usage for this research
            model_usages = (
                session.query(ModelUsage)
                .filter_by(research_id=research_id)
                .order_by(ModelUsage.total_tokens.desc())
                .all()
            )

            model_usage = []
            total_tokens = 0
            total_calls = 0

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

            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
        """
        with self.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)

            # Total tokens and researches
            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
            )

            # 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.provider)
                    .filter(ModelUsage.model_name == stat.model_name)
                    .first()
                )
                provider = (
                    provider_info.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()

            return {
                "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
                    ),
                },
            }

    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
        """
        with self.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 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,
                        "research_query": usage.research_query,
                    }
                )

            # 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,
                },
            }

    def get_research_timeline_metrics(self, research_id: int) -> 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
        """
        with self.db.get_session() 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,
            }
