"""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 ...metrics.query_utils import get_context_overflow_truncation_summary
from ...settings import SettingsManager
from ..auth.decorators import login_required

context_overflow_bp = Blueprint("context_overflow_api", __name__)

# NOTE: Routes use flask_session["username"] (not .get()) intentionally.
# @login_required guarantees the key exists; direct access fails fast
# if the decorator is ever removed.


@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["username"]

        # Get time period from query params (whitelist valid values)
        VALID_PERIODS = {"7d", "30d", "3m", "1y", "all"}
        period = request.args.get("period", "30d")
        if period not in VALID_PERIODS:
            period = "30d"

        # Pagination params for all_requests
        page = max(1, request.args.get("page", 1, type=int))
        per_page = request.args.get("per_page", 50, type=int)
        per_page = max(1, min(per_page, 500))

        # 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:
            # Truncation summary — shared with /metrics/api/metrics so the
            # main dashboard's at-a-glance numbers cannot disagree with this
            # endpoint's deep-dive. Helper internally uses
            # get_time_filter_condition, equivalent to the start_date below.
            summary = get_context_overflow_truncation_summary(session, period)
            total_requests = summary["total_requests"]
            requests_with_context = summary["requests_with_context"]
            truncated_requests = summary["truncated_requests"]
            truncation_rate = summary["truncation_rate"]
            avg_tokens_truncated = summary["avg_tokens_truncated"]

            # Base query — kept for downstream phase / chart_data / all_requests
            # aggregations that share the same time window.
            query = session.query(TokenUsage)
            if start_date:
                query = query.filter(TokenUsage.timestamp >= start_date)

            token_summary = {
                "total_requests": total_requests,
                "total_tokens": summary["total_tokens"],
                "total_prompt_tokens": summary["total_prompt_tokens"],
                "total_completion_tokens": summary["total_completion_tokens"],
                "avg_prompt_tokens": round(summary["avg_prompt_tokens"], 0),
                "avg_completion_tokens": round(
                    summary["avg_completion_tokens"], 0
                ),
                "max_prompt_tokens": summary["max_prompt_tokens"],
            }

            # --- Model token stats (always populated, no context_limit filter) ---
            model_token_query = (
                query.with_entities(
                    TokenUsage.model_name,
                    TokenUsage.model_provider,
                    func.count(TokenUsage.id).label("total_requests"),
                    func.coalesce(func.sum(TokenUsage.total_tokens), 0).label(
                        "total_tokens"
                    ),
                    func.min(TokenUsage.prompt_tokens).label("min_prompt"),
                    func.avg(TokenUsage.prompt_tokens).label("avg_prompt"),
                    func.max(TokenUsage.prompt_tokens).label("max_prompt"),
                    func.avg(TokenUsage.response_time_ms).label(
                        "avg_response_time_ms"
                    ),
                )
                .group_by(TokenUsage.model_name, TokenUsage.model_provider)
                .all()
            )

            model_token_stats = [
                {
                    "model": row.model_name,
                    "provider": row.model_provider,
                    "total_requests": row.total_requests,
                    "total_tokens": int(row.total_tokens or 0),
                    "min_prompt": int(row.min_prompt or 0),
                    "avg_prompt": round(row.avg_prompt or 0, 0),
                    "max_prompt": int(row.max_prompt or 0),
                    "avg_response_time_ms": round(
                        row.avg_response_time_ms or 0, 0
                    ),
                }
                for row in model_token_query
            ]

            # --- Phase breakdown (always populated, no context_limit filter) ---
            phase_query = (
                query.with_entities(
                    TokenUsage.research_phase,
                    func.count(TokenUsage.id).label("count"),
                    func.coalesce(func.sum(TokenUsage.total_tokens), 0).label(
                        "total_tokens"
                    ),
                    func.avg(TokenUsage.total_tokens).label("avg_tokens"),
                )
                .group_by(TokenUsage.research_phase)
                .all()
            )

            phase_breakdown = [
                {
                    "phase": row.research_phase or "unknown",
                    "count": row.count,
                    "total_tokens": int(row.total_tokens or 0),
                    "avg_tokens": round(row.avg_tokens or 0, 0),
                }
                for row in phase_query
            ]

            # 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 (capped at 1000)
                if period in ["7d", "30d"]:
                    time_series_data = time_series_query.limit(1000).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
                        "completion_tokens": usage.completion_tokens,
                        "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,
                        "provider": usage.model_provider,
                        "research_phase": usage.research_phase,
                        "response_time_ms": usage.response_time_ms,
                    }
                )

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

            # --- Paginated all_requests ---
            all_requests_query = query.order_by(desc(TokenUsage.timestamp))
            all_requests_total = all_requests_query.count()
            all_requests_pages = (
                (all_requests_total + per_page - 1) // per_page
                if all_requests_total > 0
                else 1
            )
            all_requests_data = (
                all_requests_query.offset((page - 1) * per_page)
                .limit(per_page)
                .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,
                },
                "token_summary": token_summary,
                "model_token_stats": model_token_stats,
                "phase_breakdown": phase_breakdown,
                "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,
                "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 all_requests_data
                ],
                "pagination": {
                    "page": page,
                    "per_page": per_page,
                    "total_count": all_requests_total,
                    "total_pages": all_requests_pages,
                },
                "current_context_window": SettingsManager(session).get_setting(
                    "llm.local_context_window_size"
                ),
            }

            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
