"""Zero-token task classification via ordered regex/heuristic rules.

Purely local; never calls a model. Conservative by design: when no rule
confidently matches, the task falls back to `TaskType.GENERAL` rather than
guessing, since Tier-0/Tier-1 routing decisions downstream key off this type.
"""

from __future__ import annotations

import re
from enum import StrEnum


class TaskType(StrEnum):
    ARITHMETIC = "arithmetic"
    DATE_MATH = "date_math"
    STRING_OP = "string_op"
    UNIT_CONVERSION = "unit_conversion"
    EXTRACTION = "extraction"
    CLASSIFICATION = "classification"
    MULTIPLE_CHOICE = "multiple_choice"
    SHORT_QA = "short_qa"
    CODE = "code"
    SUMMARIZATION = "summarization"
    GENERAL = "general"


_ARITHMETIC_EXPR_RE = re.compile(r"^[\s0-9+\-*/%().^]+$")
_ARITHMETIC_WORDED_RE = re.compile(
    r"\b(sum of|product of|difference|quotient|square root|percent\b|percentage|"
    r"multiplied by|divided by|\bplus\b|\bminus\b|\d+ times\b|"
    r"add \d|subtract \d|compute \d|multiply \d|"
    r"squared|cubed|raised to the power|to the power of|"
    r"marked up|discounted|\d+% tax|has a \d+% tax|markup|"
    r"how much (do you have|does it cost)|total cost|"
    r"how many .* (are needed|do you need)|"
    r"rounded to \d+ decimal places?)",
    re.IGNORECASE,
)
_ARITHMETIC_KEYWORD_RE = re.compile(r"what\s+is\s+[\d.\s+\-*/%()]+[+\-*/%]", re.IGNORECASE)
_PERCENT_OF_RE = re.compile(r"\d+(\.\d+)?\s*%\s*of\s+\d+(\.\d+)?", re.IGNORECASE)

_DATE_MATH_RE = re.compile(
    r"\b(days? (after|before|from)|weekday of|day of the week|"
    r"how many days (are there|between|until|since|from)|"
    r"is (the year|\d{4}) .*leap year|leap year|"
    r"how many days (are )?in\b|date (is|was))\b",
    re.IGNORECASE,
)
_DATE_TOKEN_RE = re.compile(
    r"\b(january|february|march|april|may|june|july|august|september|"
    r"october|november|december|\d{4}-\d{2}-\d{2}|\d{1,2}/\d{1,2}/\d{2,4})\b",
    re.IGNORECASE,
)

_STRING_OP_RE = re.compile(
    r"\b(reverse|uppercase|lowercase|upper[- ]?case|lower[- ]?case|title case|"
    r"how many (letters|characters|vowels|words)|count the (occurrences|letters)|"
    r"how many times does the letter|"
    r"first \d+ characters|last \d+ characters|sort the letters|"
    r"(\d+(st|nd|rd|th)|first|last) character of|palindrome)\b",
    re.IGNORECASE,
)

_UNIT_CONVERSION_RE = re.compile(
    r"(\bconvert\b.*\b(km|kilometers?|miles?|meters?|feet|ft|inches?|cm|"
    r"celsius|fahrenheit|kelvin|kg|kilograms?|pounds?|lbs?|grams?|"
    r"kb|mb|gb|tb|bytes?)\b|"
    r"how many (mb|kb|gb|tb|bytes?|meters?|feet|inches?|km|miles?|kg|grams?|pounds?) "
    r"(are|is) (in|there in))",
    re.IGNORECASE,
)

_EXTRACTION_RE = re.compile(
    r"\b(extract|find the (email|url|phone|number|date)s?|"
    r"list all (email|url|phone|number|date)s?)\b",
    re.IGNORECASE,
)

_MULTIPLE_CHOICE_PAREN_RE = re.compile(r"\([a-dA-D]\)", re.MULTILINE)
_MULTIPLE_CHOICE_LINE_RE = re.compile(r"^[a-dA-D][).]\s", re.MULTILINE)
# Inline "A) foo B) bar" option markers (no line anchor needed) — requires at
# least two to avoid false positives on prose containing a single "B)".
_MULTIPLE_CHOICE_INLINE_RE = re.compile(r"\b[A-D]\)\s*\S")
_MULTIPLE_CHOICE_HINT_RE = re.compile(
    r"\bchoose the correct\b|\bselect one\b|\bwhich of the following\b", re.IGNORECASE
)

_CLASSIFICATION_RE = re.compile(
    r"\b(classify|categorize|label|identify the language|"
    r"is (this|the following)\b.{0,20}\b(positive|negative|spam|ham)|sentiment)\b",
    re.IGNORECASE,
)

_CODE_RE = re.compile(
    r"\b(write a function|write code|implement|def |function\s*\(|"
    r"in (python|javascript|java|c\+\+|go|rust)\b|fix the bug|debug this|"
    r"what does this (code|python|javascript|java)?\s*(code )?(print|output|return)|"
    r"what (is|will be) the output of|what does this snippet)\b",
    re.IGNORECASE,
)

_SUMMARIZATION_RE = re.compile(
    r"\b(summarize|summary of|tl;?dr|in (one|a few) sentences?, summarize)\b",
    re.IGNORECASE,
)

_SHORT_QA_HINT_RE = re.compile(r"^(who|what|when|where|which|how many|how much)\b", re.IGNORECASE)


def classify(prompt: str) -> TaskType:
    """Classify a task prompt into a `TaskType` using ordered heuristics.

    Rules are checked most-specific-first so that, e.g., a date-math question
    phrased as a "what" question is caught by DATE_MATH before falling
    through to the generic SHORT_QA bucket. Returns GENERAL when no rule
    matches confidently.
    """
    text = prompt.strip()
    if not text:
        return TaskType.GENERAL

    stripped_expr = text.rstrip("? .")
    if _ARITHMETIC_EXPR_RE.match(stripped_expr) and any(c.isdigit() for c in text):
        return TaskType.ARITHMETIC
    if _PERCENT_OF_RE.search(text):
        return TaskType.ARITHMETIC
    if _ARITHMETIC_KEYWORD_RE.search(text) or _ARITHMETIC_WORDED_RE.search(text):
        return TaskType.ARITHMETIC

    if _DATE_MATH_RE.search(text) and _DATE_TOKEN_RE.search(text):
        return TaskType.DATE_MATH
    if _DATE_MATH_RE.search(text):
        return TaskType.DATE_MATH

    if _UNIT_CONVERSION_RE.search(text):
        return TaskType.UNIT_CONVERSION

    if _STRING_OP_RE.search(text):
        return TaskType.STRING_OP

    if _EXTRACTION_RE.search(text):
        return TaskType.EXTRACTION

    if _CODE_RE.search(text):
        return TaskType.CODE

    if _SUMMARIZATION_RE.search(text):
        return TaskType.SUMMARIZATION

    if (
        _MULTIPLE_CHOICE_PAREN_RE.search(text)
        or _MULTIPLE_CHOICE_LINE_RE.search(text)
        or len(_MULTIPLE_CHOICE_INLINE_RE.findall(text)) >= 2
        or _MULTIPLE_CHOICE_HINT_RE.search(text)
    ):
        return TaskType.MULTIPLE_CHOICE

    if _CLASSIFICATION_RE.search(text):
        return TaskType.CLASSIFICATION

    if _SHORT_QA_HINT_RE.match(text) and len(text) < 200:
        return TaskType.SHORT_QA

    return TaskType.GENERAL