from __future__ import annotations

import json
import re
from dataclasses import dataclass

from .agent_prompting import SYSTEM_PROMPT_DYNAMIC_BOUNDARY
from .agent_session import AgentMessage, AgentSessionState
from .tokenizer_runtime import describe_token_counter, count_tokens

_PATH_HEADER_RE = re.compile(r'^## ((?:/|[A-Za-z]:[\\/]).+)$', re.MULTILINE)


@dataclass(frozen=True)
class UsageEntry:
    name: str
    tokens: int


@dataclass(frozen=True)
class ToolUsageEntry:
    name: str
    call_tokens: int
    result_tokens: int


@dataclass(frozen=True)
class MessageBreakdown:
    user_message_tokens: int
    assistant_message_tokens: int
    tool_call_tokens: int
    tool_result_tokens: int
    user_context_tokens: int
    tool_calls_by_type: tuple[ToolUsageEntry, ...]


@dataclass(frozen=True)
class ContextUsageReport:
    model: str
    total_tokens: int
    raw_max_tokens: int
    percentage: float
    strategy: str
    message_count: int
    categories: tuple[UsageEntry, ...]
    system_prompt_sections: tuple[UsageEntry, ...]
    user_context_entries: tuple[UsageEntry, ...]
    system_context_entries: tuple[UsageEntry, ...]
    memory_files: tuple[UsageEntry, ...]
    message_breakdown: MessageBreakdown
    token_counter_backend: str
    token_counter_source: str
    token_counter_accurate: bool


def estimate_tokens(text: str, model: str | None = None) -> int:
    return count_tokens(text, model)


def infer_context_window(model: str) -> int:
    lowered = model.lower()
    if 'qwen3-coder' in lowered:
        return 256_000
    if 'devstral' in lowered:
        return 256_000
    if 'qwen' in lowered:
        return 131_072
    if 'claude' in lowered:
        return 200_000
    if 'gpt-4.1' in lowered or 'gpt-4o' in lowered:
        return 128_000
    return 128_000


def collect_context_usage(
    *,
    session: AgentSessionState,
    model: str,
    strategy: str,
) -> ContextUsageReport:
    raw_max_tokens = infer_context_window(model)
    token_counter = describe_token_counter(model)
    count = lambda text: estimate_tokens(text, model)  # noqa: E731
    system_prompt_sections = tuple(
        UsageEntry(name=_section_name(part, idx), tokens=count(part))
        for idx, part in enumerate(session.system_prompt_parts, start=1)
    )
    system_context_entries = tuple(
        UsageEntry(name=key, tokens=count(f'{key}: {value}'))
        for key, value in session.system_context.items()
        if value
    )
    user_context_entries = tuple(
        UsageEntry(name=key, tokens=count(_render_user_context_chunk(key, value)))
        for key, value in session.user_context.items()
        if value
    )
    memory_files = tuple(_parse_memory_usage(session.user_context.get('claudeMd'), model=model))

    user_context_tokens = sum(entry.tokens for entry in user_context_entries)
    system_prompt_tokens = (
        sum(entry.tokens for entry in system_prompt_sections)
        + sum(entry.tokens for entry in system_context_entries)
    )

    conversation_user_tokens = 0
    assistant_tokens = 0
    tool_call_tokens = 0
    tool_result_tokens = 0
    tool_usage: dict[str, list[int]] = {}

    for index, message in enumerate(session.messages):
        if index == 0 and message.role == 'system':
            continue
        if _is_user_context_message(session, index, message):
            continue
        if message.role == 'user':
            conversation_user_tokens += count(message.content)
            continue
        if message.role == 'assistant':
            assistant_tokens += count(message.content)
            for tool_call in message.tool_calls:
                serialized = json.dumps(tool_call, ensure_ascii=True)
                tokens = count(serialized)
                tool_call_tokens += tokens
                tool_name = _extract_tool_call_name(tool_call)
                call_totals = tool_usage.setdefault(tool_name, [0, 0])
                call_totals[0] += tokens
            continue
        if message.role == 'tool':
            tokens = count(message.content)
            tool_result_tokens += tokens
            result_totals = tool_usage.setdefault(message.name or 'tool', [0, 0])
            result_totals[1] += tokens

    categories = [
        UsageEntry('System prompt', system_prompt_tokens),
        UsageEntry('User context', user_context_tokens),
        UsageEntry('User messages', conversation_user_tokens),
        UsageEntry('Assistant messages', assistant_tokens),
        UsageEntry('Tool calls', tool_call_tokens),
        UsageEntry('Tool results', tool_result_tokens),
    ]
    total_tokens = sum(entry.tokens for entry in categories)
    free_space = max(raw_max_tokens - total_tokens, 0)
    categories.append(UsageEntry('Free space', free_space))

    tool_calls_by_type = tuple(
        ToolUsageEntry(
            name=name,
            call_tokens=values[0],
            result_tokens=values[1],
        )
        for name, values in sorted(
            tool_usage.items(),
            key=lambda item: (item[1][0] + item[1][1], item[0]),
            reverse=True,
        )
        if values[0] or values[1]
    )
    percentage = (total_tokens / raw_max_tokens * 100) if raw_max_tokens else 0.0
    return ContextUsageReport(
        model=model,
        total_tokens=total_tokens,
        raw_max_tokens=raw_max_tokens,
        percentage=percentage,
        strategy=strategy,
        message_count=len(session.messages),
        categories=tuple(categories),
        system_prompt_sections=system_prompt_sections,
        user_context_entries=user_context_entries,
        system_context_entries=system_context_entries,
        memory_files=memory_files,
        message_breakdown=MessageBreakdown(
            user_message_tokens=conversation_user_tokens,
            assistant_message_tokens=assistant_tokens,
            tool_call_tokens=tool_call_tokens,
            tool_result_tokens=tool_result_tokens,
            user_context_tokens=user_context_tokens,
            tool_calls_by_type=tool_calls_by_type,
        ),
        token_counter_backend=token_counter.backend,
        token_counter_source=token_counter.source,
        token_counter_accurate=token_counter.accurate,
    )


def format_context_usage(report: ContextUsageReport) -> str:
    lines = [
        '## Context Usage',
        '',
        f'**Model:** {report.model}  ',
        f'**Estimated tokens:** {_format_tokens(report.total_tokens)} / {_format_tokens(report.raw_max_tokens)} ({report.percentage:.1f}%)  ',
        f'**Token counter:** {report.token_counter_backend} ({report.token_counter_source}){" [accurate]" if report.token_counter_accurate else " [fallback]"}  ',
        f'**Context strategy:** {report.strategy}  ',
        f'**Messages in session:** {report.message_count}',
        '',
    ]

    visible_categories = [entry for entry in report.categories if entry.tokens > 0]
    if visible_categories:
        lines.extend(
            [
                '### Estimated usage by category',
                '',
                '| Category | Tokens | Percentage |',
                '|----------|--------|------------|',
            ]
        )
        for entry in visible_categories:
            percent = (entry.tokens / report.raw_max_tokens * 100) if report.raw_max_tokens else 0.0
            lines.append(f'| {entry.name} | {_format_tokens(entry.tokens)} | {percent:.1f}% |')
        lines.append('')

    if report.system_prompt_sections:
        lines.extend(
            [
                '### System Prompt Sections',
                '',
                '| Section | Tokens |',
                '|---------|--------|',
            ]
        )
        for entry in report.system_prompt_sections:
            lines.append(f'| {entry.name} | {_format_tokens(entry.tokens)} |')
        lines.append('')

    if report.user_context_entries:
        lines.extend(
            [
                '### User Context',
                '',
                '| Entry | Tokens |',
                '|-------|--------|',
            ]
        )
        for entry in report.user_context_entries:
            lines.append(f'| {entry.name} | {_format_tokens(entry.tokens)} |')
        lines.append('')

    if report.system_context_entries:
        lines.extend(
            [
                '### System Context',
                '',
                '| Entry | Tokens |',
                '|-------|--------|',
            ]
        )
        for entry in report.system_context_entries:
            lines.append(f'| {entry.name} | {_format_tokens(entry.tokens)} |')
        lines.append('')

    if report.memory_files:
        lines.extend(
            [
                '### Memory Files',
                '',
                '| Path | Tokens |',
                '|------|--------|',
            ]
        )
        for entry in report.memory_files:
            lines.append(f'| {entry.name} | {_format_tokens(entry.tokens)} |')
        lines.append('')

    breakdown = report.message_breakdown
    lines.extend(
        [
            '### Message Breakdown',
            '',
            '| Category | Tokens |',
            '|----------|--------|',
            f'| User context reminder | {_format_tokens(breakdown.user_context_tokens)} |',
            f'| User messages | {_format_tokens(breakdown.user_message_tokens)} |',
            f'| Assistant messages | {_format_tokens(breakdown.assistant_message_tokens)} |',
            f'| Tool calls | {_format_tokens(breakdown.tool_call_tokens)} |',
            f'| Tool results | {_format_tokens(breakdown.tool_result_tokens)} |',
            '',
        ]
    )

    if breakdown.tool_calls_by_type:
        lines.extend(
            [
                '#### Top Tools',
                '',
                '| Tool | Call Tokens | Result Tokens |',
                '|------|-------------|---------------|',
            ]
        )
        for entry in breakdown.tool_calls_by_type:
            lines.append(
                f'| {entry.name} | {_format_tokens(entry.call_tokens)} | {_format_tokens(entry.result_tokens)} |'
            )
        lines.append('')

    while lines and lines[-1] == '':
        lines.pop()
    return '\n'.join(lines)


def _section_name(part: str, index: int) -> str:
    stripped = part.strip()
    if stripped == SYSTEM_PROMPT_DYNAMIC_BOUNDARY:
        return 'Dynamic boundary'
    first_line = stripped.splitlines()[0] if stripped else ''
    if first_line.startswith('#'):
        return first_line.lstrip('#').strip() or f'Section {index}'
    return f'Section {index}'


def _render_user_context_chunk(key: str, value: str) -> str:
    return f'# {key}\n{value}'


def _extract_tool_call_name(tool_call: dict[str, object]) -> str:
    function_block = tool_call.get('function')
    if isinstance(function_block, dict):
        name = function_block.get('name')
        if isinstance(name, str) and name:
            return name
    return 'unknown'


def _is_user_context_message(
    session: AgentSessionState,
    index: int,
    message: AgentMessage,
) -> bool:
    if not session.user_context:
        return False
    return (
        index == 1
        and message.role == 'user'
        and message.content.startswith('<system-reminder>')
    )


def _parse_memory_usage(claude_md: str | None, *, model: str | None = None) -> list[UsageEntry]:
    if not claude_md:
        return []
    matches = list(_PATH_HEADER_RE.finditer(claude_md))
    if not matches:
        return []
    entries: list[UsageEntry] = []
    for idx, match in enumerate(matches):
        start = match.end()
        end = matches[idx + 1].start() if idx + 1 < len(matches) else len(claude_md)
        content = claude_md[start:end].strip()
        entries.append(UsageEntry(name=match.group(1), tokens=estimate_tokens(content, model)))
    return entries


def _format_tokens(value: int) -> str:
    return f'{value:,}'