# -*- coding: utf-8 -*-
"""Agent statistics service."""

from __future__ import annotations

import asyncio
import logging
import os
from datetime import date, timedelta
from pathlib import Path

import aiofiles
import aiofiles.os
import orjson

from ..app.chats.repo import JsonChatRepository
from ..token_usage import get_token_usage_manager
from .models import (
    AgentStatsSummary,
    ChannelStats,
    DailyStats,
)

logger = logging.getLogger(__name__)


# pylint: disable=unused-argument
def _should_skip_by_mtime(
    session_file: Path,
    start_date: date,
    end_date: date,
) -> bool:
    try:
        mtime = session_file.stat().st_mtime
        mtime_date = date.fromtimestamp(mtime)
        if mtime_date < start_date:
            logger.debug(
                "Skipping %s by mtime (%s) before start date %s",
                session_file.name,
                mtime_date.isoformat(),
                start_date.isoformat(),
            )
            return True
    except OSError:
        pass
    return False


def _extract_session_messages(session_data: dict) -> list:
    """Return raw message dicts/tuples from a session state, 1.x or 2.0."""
    agent_raw = session_data.get("agent", {})
    # 2.0: messages live on agent.state.context
    state_raw = agent_raw.get("state")
    if isinstance(state_raw, dict):
        ctx = state_raw.get("context")
        if isinstance(ctx, list) and ctx:
            return ctx
    # 1.x fallback
    memory_raw = agent_raw.get("memory", {})
    if isinstance(memory_raw, dict):
        return memory_raw.get("memories") or memory_raw.get("content") or []
    return []


def _should_skip_by_content_range(
    session_data: dict,
    start_date_str: str,
    end_date_str: str,
) -> bool:
    memories = _extract_session_messages(session_data)

    if not memories:
        return True

    timestamps: list[str] = []
    for msg_item in memories:
        if isinstance(msg_item, list) and len(msg_item) > 0:
            msg_data = msg_item[0]
        elif isinstance(msg_item, dict):
            msg_data = msg_item
        else:
            continue

        if not isinstance(msg_data, dict):
            continue

        timestamp = msg_data.get("created_at") or msg_data.get("timestamp")
        if timestamp:
            timestamps.append(str(timestamp)[:10])

    if not timestamps:
        return True

    first_date = timestamps[0]
    last_date = timestamps[-1]

    if last_date < start_date_str or first_date > end_date_str:
        logger.debug(
            "Skipping session by content range [%s, %s] "
            "outside target [%s, %s]",
            first_date,
            last_date,
            start_date_str,
            end_date_str,
        )
        return True

    return False


# pylint:disable=too-many-statements,too-many-branches
def _process_session_file(
    session_data: dict,
    start_date_str: str,
    end_date_str: str,
    daily_stats: dict[str, dict],
    channel_stats: dict[str, dict],
    channel: str,
    session_stem: str,
    active_sessions: dict[str, set[str]],
) -> tuple[int, bool]:
    tool_call_count = 0
    has_messages_in_range = False
    try:
        memories = _extract_session_messages(session_data)

        stats = channel_stats.setdefault(
            channel,
            {
                "session_count": 0,
                "user_messages": 0,
                "assistant_messages": 0,
                "total_messages": 0,
            },
        )

        for msg_item in memories:
            if isinstance(msg_item, list) and len(msg_item) > 0:
                msg_data = msg_item[0]
            elif isinstance(msg_item, dict):
                msg_data = msg_item
            else:
                continue

            if not isinstance(msg_data, dict):
                continue

            timestamp = msg_data.get("created_at") or msg_data.get("timestamp")
            if not timestamp:
                continue

            date_str = str(timestamp)[:10]
            if date_str < start_date_str or date_str > end_date_str:
                continue

            has_messages_in_range = True
            active_sessions.setdefault(date_str, set()).add(session_stem)

            ds = daily_stats[date_str]
            role = msg_data.get("role", "")
            content = msg_data.get("content", [])

            if role == "user":
                ds["user_messages"] += 1
                ds["total_messages"] += 1
                stats["user_messages"] += 1
                stats["total_messages"] += 1
            elif role == "assistant":
                ds["assistant_messages"] += 1
                ds["total_messages"] += 1
                stats["assistant_messages"] += 1
                stats["total_messages"] += 1

            if isinstance(content, list):
                for block in content:
                    btype = (
                        block.get("type")
                        if isinstance(block, dict)
                        else getattr(block, "type", None)
                    )
                    if btype in ("tool_use", "tool_call"):
                        ds["tool_calls"] += 1
                        tool_call_count += 1

    except Exception as e:
        logger.debug("Failed to count messages in session: %s", e)

    if has_messages_in_range and channel in channel_stats:
        channel_stats[channel]["session_count"] += 1

    return tool_call_count, has_messages_in_range


class AgentStatsService:
    """Service for computing agent statistics."""

    # pylint: disable=R0912,R0915
    async def get_summary(
        self,
        workspace_dir: Path,
        start_date: date,
        end_date: date,
    ) -> AgentStatsSummary:
        chats_file = workspace_dir / "chats.json"
        sessions_dir = workspace_dir / "sessions"

        daily_stats: dict[str, dict] = {}
        days = (end_date - start_date).days + 1
        for i in range(days):
            date_str = (start_date + timedelta(days=i)).isoformat()
            daily_stats[date_str] = {
                "date": date_str,
                "chats": 0,
                "active_sessions": 0,
                "user_messages": 0,
                "assistant_messages": 0,
                "total_messages": 0,
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "llm_calls": 0,
                "tool_calls": 0,
            }

        start_date_str = start_date.isoformat()
        end_date_str = end_date.isoformat()

        channel_stats: dict[str, dict] = {}
        total_tool_calls = 0
        active_sessions: dict[str, set[str]] = {}
        total_active_sessions = 0

        if chats_file.exists():
            try:
                repo = JsonChatRepository(chats_file)
                chats = await repo.list_chats()
                for chat in chats:
                    if chat.created_at is None:
                        continue
                    chat_date = chat.created_at.date()
                    if start_date <= chat_date <= end_date:
                        date_str = chat_date.isoformat()
                        daily_stats[date_str]["chats"] += 1
            except Exception as e:
                logger.warning("Failed to load chat statistics: %s", e)

        # pylint: disable=too-many-nested-blocks
        if sessions_dir.exists():
            try:
                session_files = []

                # Scan root sessions directory for legacy files
                channel_names = await aiofiles.os.listdir(sessions_dir)
                for channel_name in channel_names:
                    channel_path = sessions_dir / channel_name
                    if await aiofiles.os.path.isdir(channel_path):
                        try:
                            channel_files = await aiofiles.os.listdir(
                                channel_path,
                            )
                            for channel_file in channel_files:
                                session_file = channel_path / channel_file
                                if session_file.name.endswith(".json"):
                                    session_files.append(session_file)
                        except Exception as e:
                            logger.debug(
                                "Failed to scan channel directory %s: %s",
                                channel_path,
                                e,
                            )

                session_fd_sem = asyncio.Semaphore((os.cpu_count() or 4) * 2)

                async def _process_one(session_file: Path) -> tuple[int, bool]:
                    async with session_fd_sem:
                        if _should_skip_by_mtime(
                            session_file,
                            start_date,
                            end_date,
                        ):
                            return 0, False

                        try:
                            async with aiofiles.open(
                                session_file,
                                "r",
                                encoding="utf-8",
                            ) as f:
                                session_data = orjson.loads(await f.read())
                        except Exception as e:
                            logger.debug(
                                "Failed to read session file %s: %s",
                                session_file,
                                e,
                            )
                            return 0, False

                        if _should_skip_by_content_range(
                            session_data,
                            start_date_str,
                            end_date_str,
                        ):
                            return 0, False

                        stem = session_file.stem
                        # Check if session is in a channel subdirectory
                        channel = session_file.parent.name

                        return _process_session_file(
                            session_data,
                            start_date_str,
                            end_date_str,
                            daily_stats,
                            channel_stats,
                            channel,
                            stem,
                            active_sessions,
                        )

                tasks = [_process_one(sf) for sf in session_files]
                results = await asyncio.gather(*tasks, return_exceptions=True)
                for result in results:
                    if isinstance(result, tuple) and len(result) == 2:
                        tool_calls, has_messages = result
                        total_tool_calls += tool_calls
                        if has_messages:
                            total_active_sessions += 1
                    elif isinstance(result, Exception):
                        logger.debug("Failed to process session: %s", result)
            except Exception as e:
                logger.warning("Failed to load message statistics: %s", e)

        token_summary = await get_token_usage_manager().get_summary(
            start_date=start_date,
            end_date=end_date,
        )
        for date_str, ts in token_summary.by_date.items():
            if date_str in daily_stats:
                daily_stats[date_str]["prompt_tokens"] = ts.prompt_tokens
                daily_stats[date_str][
                    "completion_tokens"
                ] = ts.completion_tokens
                daily_stats[date_str]["llm_calls"] = ts.call_count

        for date_str, session_set in active_sessions.items():
            if date_str in daily_stats:
                daily_stats[date_str]["active_sessions"] = len(session_set)

        by_date = [daily_stats[d] for d in sorted(daily_stats.keys())]

        total_user_messages = sum(ds["user_messages"] for ds in by_date)
        total_assistant_messages = sum(
            ds["assistant_messages"] for ds in by_date
        )
        total_messages = total_user_messages + total_assistant_messages

        return AgentStatsSummary(
            total_active_sessions=total_active_sessions,
            total_messages=total_messages,
            total_user_messages=total_user_messages,
            total_assistant_messages=total_assistant_messages,
            total_prompt_tokens=token_summary.total_prompt_tokens,
            total_completion_tokens=token_summary.total_completion_tokens,
            total_llm_calls=token_summary.total_calls,
            total_tool_calls=total_tool_calls,
            by_date=[DailyStats.model_validate(ds) for ds in by_date],
            channel_stats=[
                ChannelStats(
                    channel=ch,
                    session_count=cnts["session_count"],
                    user_messages=cnts["user_messages"],
                    assistant_messages=cnts["assistant_messages"],
                    total_messages=cnts["total_messages"],
                )
                for ch, cnts in sorted(channel_stats.items())
            ],
            start_date=start_date_str,
            end_date=end_date_str,
        )


_agent_stats_service: AgentStatsService | None = None


def get_agent_stats_service() -> AgentStatsService:
    global _agent_stats_service
    if _agent_stats_service is None:
        _agent_stats_service = AgentStatsService()
    return _agent_stats_service
