"""
Task classifier — determines prompt complexity locally (no network call).

Uses heuristics based on prompt length, structure, and keyword analysis.
This runs entirely on the client side for zero latency.
"""

import re
from costwise_mcp.models import TaskComplexity


# Keywords that indicate complex tasks
_COMPLEX_KEYWORDS = frozenset([
    "analyze", "compare", "evaluate", "synthesize", "design", "architect",
    "implement", "refactor", "debug", "optimize", "explain in detail",
    "write a full", "create a complete", "build", "develop", "generate code",
    "step by step", "comprehensive", "thorough", "in-depth", "detailed analysis",
    "multi-step", "pipeline", "workflow", "system design", "trade-offs",
    "pros and cons", "research", "literature review", "essay", "report",
])

# Keywords that indicate simple tasks
_SIMPLE_KEYWORDS = frozenset([
    "translate", "summarize", "tldr", "fix grammar", "correct spelling",
    "yes or no", "true or false", "one word", "short answer", "brief",
    "classify", "label", "tag", "extract", "list", "name", "define",
    "what is", "who is", "convert", "format", "rewrite briefly",
])


def classify_task(prompt: str) -> TaskComplexity:
    """
    Classify task complexity from the prompt text.

    Rules (applied in order):
    1. Very short prompts (< 50 chars) → SIMPLE
    2. Keyword matching for SIMPLE/COMPLEX indicators
    3. Prompt length heuristic
    4. Structural complexity (multiple questions, code blocks, lists)

    Returns:
        TaskComplexity.SIMPLE, MEDIUM, or COMPLEX
    """
    if not prompt or not prompt.strip():
        return TaskComplexity.SIMPLE

    prompt_lower = prompt.lower().strip()
    prompt_len = len(prompt)
    word_count = len(prompt.split())

    # ── Rule 1: Very short prompts ────────────────────────────────
    if prompt_len < 50 or word_count < 10:
        return TaskComplexity.SIMPLE

    # ── Rule 2: Keyword matching ──────────────────────────────────
    simple_score = sum(1 for kw in _SIMPLE_KEYWORDS if kw in prompt_lower)
    complex_score = sum(1 for kw in _COMPLEX_KEYWORDS if kw in prompt_lower)

    if simple_score > 0 and complex_score == 0:
        return TaskComplexity.SIMPLE
    if complex_score >= 2:
        return TaskComplexity.COMPLEX

    # ── Rule 3: Structural complexity ─────────────────────────────
    # Multiple questions
    question_count = prompt.count("?")
    # Code blocks
    code_blocks = len(re.findall(r"```", prompt))
    # Numbered lists
    numbered_items = len(re.findall(r"^\s*\d+[\.\)]\s", prompt, re.MULTILINE))
    # Bullet points
    bullet_items = len(re.findall(r"^\s*[-*]\s", prompt, re.MULTILINE))

    structural_score = question_count + code_blocks + numbered_items + bullet_items

    if structural_score >= 4:
        return TaskComplexity.COMPLEX
    if structural_score >= 2:
        return TaskComplexity.MEDIUM

    # ── Rule 4: Length heuristic ──────────────────────────────────
    if word_count > 500 or prompt_len > 2000:
        return TaskComplexity.COMPLEX
    if word_count > 100 or prompt_len > 500:
        return TaskComplexity.MEDIUM

    return TaskComplexity.MEDIUM


def estimate_tokens(text: str) -> int:
    """
    Fast token estimation without loading tiktoken (for speed).
    Uses the ~4 chars per token heuristic for English text.
    Accurate to within ~10% for most prompts.
    """
    if not text:
        return 0
    # Simple heuristic: ~4 chars per token for English
    # Slightly more accurate: count words * 1.3
    word_count = len(text.split())
    char_estimate = len(text) // 4
    word_estimate = int(word_count * 1.3)
    return max(char_estimate, word_estimate)


def estimate_tokens_precise(text: str, model: str = "gpt-4") -> int:
    """
    Precise token count using tiktoken (slower, requires tiktoken installed).
    Falls back to heuristic if tiktoken is not available.
    """
    try:
        import tiktoken
        # Map model names to tiktoken encodings
        if "gpt-4" in model or "gpt-3.5" in model:
            enc = tiktoken.encoding_for_model(model)
        else:
            enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except (ImportError, KeyError):
        return estimate_tokens(text)
