"""Token counting and grading logic for MCP tool definitions."""

from __future__ import annotations

import json
from dataclasses import dataclass

import tiktoken


GRADING_THRESHOLDS = [
    (200, "A"),
    (500, "B"),
    (1000, "C"),
    (2000, "D"),
]

DEFAULT_BUDGET = 200_000

# Anthropic Claude pricing per million tokens (input)
# Claude 3.5 Sonnet: $3/M input
PRICE_PER_MILLION_INPUT = 3.0

_encoder: tiktoken.Encoding | None = None


def _get_encoder() -> tiktoken.Encoding:
    global _encoder
    if _encoder is None:
        _encoder = tiktoken.get_encoding("cl100k_base")
    return _encoder


def count_tokens(text: str) -> int:
    """Count tokens using cl100k_base encoding."""
    return len(_get_encoder().encode(text))


def tool_to_text(tool) -> str:
    """Serialize a tool definition the way it would appear in an API call."""
    parts = [f"Tool: {tool.name}"]
    if tool.description:
        parts.append(f"Description: {tool.description}")
    if tool.inputSchema:
        parts.append(f"Input Schema: {json.dumps(tool.inputSchema, separators=(',', ':'))}")
    return "\n".join(parts)


def schema_depth(schema: dict | None, current: int = 0) -> int:
    """Calculate maximum nesting depth of a JSON schema."""
    if not schema or not isinstance(schema, dict):
        return current
    max_depth = current
    properties = schema.get("properties", {})
    if properties:
        for prop in properties.values():
            if isinstance(prop, dict):
                max_depth = max(max_depth, schema_depth(prop, current + 1))
    items = schema.get("items")
    if isinstance(items, dict):
        max_depth = max(max_depth, schema_depth(items, current + 1))
    for key in ("allOf", "anyOf", "oneOf"):
        variants = schema.get(key)
        if isinstance(variants, list):
            for variant in variants:
                if isinstance(variant, dict):
                    max_depth = max(max_depth, schema_depth(variant, current + 1))
    return max_depth


def grade_tokens(tokens: int) -> str:
    """Assign A-F grade based on token count."""
    for threshold, g in GRADING_THRESHOLDS:
        if tokens < threshold:
            return g
    return "F"


@dataclass
class ToolAnalysis:
    name: str
    tokens: int
    budget_pct: float
    schema_depth: int
    desc_length: int
    grade: str
    description: str | None = None
    input_schema: dict | None = None


def analyze_tool(tool, budget: int = DEFAULT_BUDGET) -> ToolAnalysis:
    """Analyze a single MCP tool definition."""
    text = tool_to_text(tool)
    tokens = count_tokens(text)
    depth = schema_depth(tool.inputSchema)
    desc_len = len(tool.description) if tool.description else 0
    return ToolAnalysis(
        name=tool.name,
        tokens=tokens,
        budget_pct=(tokens / budget) * 100,
        schema_depth=depth,
        desc_length=desc_len,
        grade=grade_tokens(tokens),
        description=tool.description,
        input_schema=tool.inputSchema,
    )


@dataclass
class AnalysisSummary:
    total_tokens: int
    budget: int
    budget_pct: float
    top_offenders: list[ToolAnalysis]
    cost_per_1000: float
    tools: list[ToolAnalysis]


def analyze_tools(tools, budget: int = DEFAULT_BUDGET) -> AnalysisSummary:
    """Analyze all tools from an MCP server."""
    analyses = [analyze_tool(t, budget) for t in tools]
    total = sum(a.tokens for a in analyses)
    sorted_by_tokens = sorted(analyses, key=lambda a: a.tokens, reverse=True)
    cost = (total / 1_000_000) * PRICE_PER_MILLION_INPUT * 1000
    return AnalysisSummary(
        total_tokens=total,
        budget=budget,
        budget_pct=(total / budget) * 100,
        top_offenders=sorted_by_tokens[:3],
        cost_per_1000=cost,
        tools=analyses,
    )
