"""API endpoints for context overflow analytics."""

from flask import Blueprint, jsonify, request, session as flask_session
from datetime import datetime, timedelta, timezone
from sqlalchemy import func, desc
from loguru import logger

from ...database.session_context import get_user_db_session
from ...database.models import TokenUsage
from ..auth.decorators import login_required

context_overflow_bp = Blueprint("context_overflow_api", __name__)


@context_overflow_bp.route("/api/context-overflow", methods=["GET"])
@login_required
def get_context_overflow_metrics():
    """Get context overflow metrics for the current user."""
    try:
        # Get username from session
        username = flask_session.get("username")
        if not username:
            return jsonify(
                {"status": "error", "message": "User not authenticated"}
            ), 401

        # Get time period from query params
        period = request.args.get("period", "30d")

        # Calculate date filter (use timezone-aware datetime)
        start_date = None
        if period != "all":
            now = datetime.now(timezone.utc)
            if period == "7d":
                start_date = now - timedelta(days=7)
            elif period == "30d":
                start_date = now - timedelta(days=30)
            elif period == "3m":
                start_date = now - timedelta(days=90)
            elif period == "1y":
                start_date = now - timedelta(days=365)

        with get_user_db_session(username) as session:
            # Base query
            query = session.query(TokenUsage)

            if start_date:
                query = query.filter(TokenUsage.timestamp >= start_date)

            # Get overview statistics
            total_requests = query.count()

            # Requests with context data
            requests_with_context = query.filter(
                TokenUsage.context_limit.isnot(None)
            ).count()

            # Truncated requests
            truncated_requests = query.filter(
                TokenUsage.context_truncated.is_(True)
            ).count()

            # Calculate truncation rate
            truncation_rate = 0
            if requests_with_context > 0:
                truncation_rate = (
                    truncated_requests / requests_with_context
                ) * 100

            # Get average tokens truncated
            avg_tokens_truncated = session.query(
                func.avg(TokenUsage.tokens_truncated)
            ).filter(TokenUsage.context_truncated.is_(True))

            if start_date:
                avg_tokens_truncated = avg_tokens_truncated.filter(
                    TokenUsage.timestamp >= start_date
                )

            avg_tokens_truncated = avg_tokens_truncated.scalar() or 0

            # Get context limit distribution by model
            context_limits = session.query(
                TokenUsage.model_name,
                TokenUsage.context_limit,
                func.count(TokenUsage.id).label("count"),
            ).filter(TokenUsage.context_limit.isnot(None))

            if start_date:
                context_limits = context_limits.filter(
                    TokenUsage.timestamp >= start_date
                )

            context_limits = context_limits.group_by(
                TokenUsage.model_name, TokenUsage.context_limit
            ).all()

            # Get recent truncated requests
            recent_truncated = (
                query.filter(TokenUsage.context_truncated.is_(True))
                .order_by(desc(TokenUsage.timestamp))
                .limit(20)
                .all()
            )

            # Get time series data for chart - include all records
            # (even those without context_limit for OpenRouter models)
            time_series_query = query.order_by(TokenUsage.timestamp)

            if start_date:
                # For shorter periods, get all data points
                if period in ["7d", "30d"]:
                    time_series_data = time_series_query.all()
                else:
                    # For longer periods, sample data
                    time_series_data = time_series_query.limit(500).all()
            else:
                time_series_data = time_series_query.limit(1000).all()

            # Format time series for chart
            chart_data = []
            for usage in time_series_data:
                # Calculate original tokens (before truncation)
                ollama_used = (
                    usage.ollama_prompt_eval_count
                )  # What Ollama actually processed
                actual_prompt = ollama_used or usage.prompt_tokens
                tokens_truncated = usage.tokens_truncated or 0
                original_tokens = (
                    actual_prompt + tokens_truncated
                    if usage.context_truncated
                    else actual_prompt
                )

                chart_data.append(
                    {
                        "timestamp": usage.timestamp.isoformat(),
                        "research_id": usage.research_id,
                        "prompt_tokens": usage.prompt_tokens,  # From our standard token counting
                        "ollama_prompt_tokens": ollama_used,  # What Ollama actually used (may be capped)
                        "original_prompt_tokens": original_tokens,  # What was originally requested (before truncation)
                        "context_limit": usage.context_limit,
                        "truncated": bool(usage.context_truncated),
                        "tokens_truncated": tokens_truncated,
                        "model": usage.model_name,
                    }
                )

            # Get model-specific truncation stats
            model_stats = session.query(
                TokenUsage.model_name,
                TokenUsage.model_provider,
                func.count(TokenUsage.id).label("total_requests"),
                func.sum(TokenUsage.context_truncated).label("truncated_count"),
                func.avg(TokenUsage.context_limit).label("avg_context_limit"),
            ).filter(TokenUsage.context_limit.isnot(None))

            if start_date:
                model_stats = model_stats.filter(
                    TokenUsage.timestamp >= start_date
                )

            model_stats = model_stats.group_by(
                TokenUsage.model_name, TokenUsage.model_provider
            ).all()

            # Format response
            response = {
                "status": "success",
                "overview": {
                    "total_requests": total_requests,
                    "requests_with_context_data": requests_with_context,
                    "truncated_requests": truncated_requests,
                    "truncation_rate": round(truncation_rate, 2),
                    "avg_tokens_truncated": round(avg_tokens_truncated, 0)
                    if avg_tokens_truncated
                    else 0,
                },
                "context_limits": [
                    {"model": model, "limit": limit, "count": count}
                    for model, limit, count in context_limits
                ],
                "model_stats": [
                    {
                        "model": stat.model_name,
                        "provider": stat.model_provider,
                        "total_requests": stat.total_requests,
                        "truncated_count": int(stat.truncated_count or 0),
                        "truncation_rate": round(
                            (stat.truncated_count or 0)
                            / stat.total_requests
                            * 100,
                            2,
                        )
                        if stat.total_requests > 0
                        else 0,
                        "avg_context_limit": round(stat.avg_context_limit, 0)
                        if stat.avg_context_limit
                        else None,
                    }
                    for stat in model_stats
                ],
                "recent_truncated": [
                    {
                        "timestamp": req.timestamp.isoformat(),
                        "research_id": req.research_id,
                        "model": req.model_name,
                        "prompt_tokens": req.prompt_tokens,  # Standard token count
                        "ollama_tokens": req.ollama_prompt_eval_count,  # What Ollama actually used
                        "original_tokens": (
                            req.ollama_prompt_eval_count or req.prompt_tokens
                        )
                        + (req.tokens_truncated or 0),  # What was requested
                        "context_limit": req.context_limit,
                        "tokens_truncated": req.tokens_truncated,
                        "truncation_ratio": req.truncation_ratio,
                        "research_query": req.research_query,
                    }
                    for req in recent_truncated
                ],
                "chart_data": chart_data,
                # Add detailed table data for all requests
                "all_requests": [
                    {
                        "timestamp": req.timestamp.isoformat(),
                        "research_id": req.research_id,
                        "model": req.model_name,
                        "provider": req.model_provider,
                        "prompt_tokens": req.prompt_tokens,
                        "completion_tokens": req.completion_tokens,
                        "total_tokens": req.total_tokens,
                        "context_limit": req.context_limit,
                        "context_truncated": bool(req.context_truncated),
                        "tokens_truncated": req.tokens_truncated or 0,
                        "truncation_ratio": round(req.truncation_ratio * 100, 2)
                        if req.truncation_ratio
                        else 0,
                        "ollama_prompt_eval_count": req.ollama_prompt_eval_count,
                        "research_query": req.research_query,
                        "research_phase": req.research_phase,
                    }
                    for req in query.order_by(desc(TokenUsage.timestamp))
                    .limit(100)
                    .all()
                ],
            }

            return jsonify(response)

    except Exception:
        logger.exception("Error getting context overflow metrics")
        return jsonify(
            {
                "status": "error",
                "message": "Failed to load context overflow metrics",
            }
        ), 500


@context_overflow_bp.route(
    "/api/research/<string:research_id>/context-overflow", methods=["GET"]
)
@login_required
def get_research_context_overflow(research_id):
    """Get context overflow metrics for a specific research."""
    try:
        with get_user_db_session() as session:
            # Get all token usage for this research
            token_usage = (
                session.query(TokenUsage)
                .filter(TokenUsage.research_id == research_id)
                .order_by(TokenUsage.timestamp)
                .all()
            )

            if not token_usage:
                return jsonify(
                    {
                        "status": "success",
                        "data": {
                            "overview": {
                                "total_requests": 0,
                                "total_tokens": 0,
                                "context_limit": None,
                                "max_tokens_used": 0,
                                "truncation_occurred": False,
                            },
                            "requests": [],
                        },
                    }
                )

            # Calculate overview metrics
            total_tokens = sum(req.total_tokens or 0 for req in token_usage)
            total_prompt = sum(req.prompt_tokens or 0 for req in token_usage)
            total_completion = sum(
                req.completion_tokens or 0 for req in token_usage
            )

            # Get context limit (should be same for all requests in a research)
            context_limit = next(
                (req.context_limit for req in token_usage if req.context_limit),
                None,
            )

            # Check for truncation
            truncated_requests = [
                req for req in token_usage if req.context_truncated
            ]
            max_tokens_used = max(
                (req.prompt_tokens or 0) for req in token_usage
            )

            # Get token usage by phase
            phase_stats = {}
            for req in token_usage:
                phase = req.research_phase or "unknown"
                if phase not in phase_stats:
                    phase_stats[phase] = {
                        "count": 0,
                        "prompt_tokens": 0,
                        "completion_tokens": 0,
                        "total_tokens": 0,
                        "truncated_count": 0,
                    }
                phase_stats[phase]["count"] += 1
                phase_stats[phase]["prompt_tokens"] += req.prompt_tokens or 0
                phase_stats[phase]["completion_tokens"] += (
                    req.completion_tokens or 0
                )
                phase_stats[phase]["total_tokens"] += req.total_tokens or 0
                if req.context_truncated:
                    phase_stats[phase]["truncated_count"] += 1

            # Format requests for response
            requests_data = []
            for req in token_usage:
                requests_data.append(
                    {
                        "timestamp": req.timestamp.isoformat(),
                        "phase": req.research_phase,
                        "prompt_tokens": req.prompt_tokens,
                        "completion_tokens": req.completion_tokens,
                        "total_tokens": req.total_tokens,
                        "context_limit": req.context_limit,
                        "context_truncated": bool(req.context_truncated),
                        "tokens_truncated": req.tokens_truncated or 0,
                        "ollama_prompt_eval_count": req.ollama_prompt_eval_count,
                        "calling_function": req.calling_function,
                        "response_time_ms": req.response_time_ms,
                    }
                )

            response = {
                "status": "success",
                "data": {
                    "overview": {
                        "total_requests": len(token_usage),
                        "total_tokens": total_tokens,
                        "total_prompt_tokens": total_prompt,
                        "total_completion_tokens": total_completion,
                        "context_limit": context_limit,
                        "max_tokens_used": max_tokens_used,
                        "truncation_occurred": len(truncated_requests) > 0,
                        "truncated_count": len(truncated_requests),
                        "tokens_lost": sum(
                            req.tokens_truncated or 0
                            for req in truncated_requests
                        ),
                    },
                    "phase_stats": phase_stats,
                    "requests": requests_data,
                    "model": token_usage[0].model_name if token_usage else None,
                    "provider": token_usage[0].model_provider
                    if token_usage
                    else None,
                },
            }

            return jsonify(response)

    except Exception:
        logger.exception("Error getting research context overflow")
        return jsonify(
            {
                "status": "error",
                "message": "Failed to load context overflow data",
            }
        ), 500
