from __future__ import annotations
from typing import Dict, List, Tuple
from smolagents import tool

# Import only the classifier API; DO NOT construct models here.
from level_classifier_tool_2 import classify_levels_phrases
from phrases import BLOOMS_PHRASES, DOK_PHRASES

# ------------------------ Injected state (set from app.py) ------------------------
_INDEX = None
_BACKEND = None
_BLOOM_INDEX = None
_DOK_INDEX = None

def set_retrieval_index(index) -> None:
    """Call this from app.py after loading your LlamaIndex index."""
    global _INDEX
    _INDEX = index

def set_classifier_state(backend, bloom_index, dok_index) -> None:
    """Call this from app.py after building the backend and prebuilt indices."""
    global _BACKEND, _BLOOM_INDEX, _DOK_INDEX
    _BACKEND = backend
    _BLOOM_INDEX = bloom_index
    _DOK_INDEX = dok_index

# ----------------------------- Tools -------------------------------------

@tool
def QuestionRetrieverTool(subject: str, topic: str, grade: str) -> dict:
    """
    Retrieve up to 5 closely-related example Q&A pairs from the source datasets.

    Args:
        subject: The subject area (e.g., "Math", "Science").
        topic: The specific topic within the subject (e.g., "Algebra", "Biology").
        grade: The grade level (e.g., "Grade 5", "Grade 8").

    Returns:
        {
          "closest questions found for": {"subject": ..., "topic": ..., "grade": ...},
          "questions": [{"text": "..."} * up to 5]
        }
    """
    if _INDEX is None:
        return {"error": "Retriever not initialized. Call set_retrieval_index(index) before using this tool."}

    query = f"{topic} question for {grade} of the {subject}"
    try:
        results = _INDEX.as_retriever(similarity_top_k=5).retrieve(query)
        question_texts = [r.node.text for r in results]
    except Exception as e:
        return {"error": f"Retriever error: {e}"}

    return {
        "closest questions found for": {"subject": subject, "topic": topic, "grade": grade},
        "questions": [{"text": q} for q in question_texts]
    }


@tool
def classify_and_score(
    question: str,
    target_bloom: str,
    target_dok: str,
    agg: str = "max"
) -> dict:
    """
    Classify a question against Bloom’s and DOK targets and return guidance.

    Args:
        question: Question text to evaluate.
        target_bloom: Target Bloom’s level (e.g., "Analyze" or "Apply+").
        target_dok: Target DOK level (e.g., "DOK3" or "DOK2-DOK3").
        agg: Aggregation over phrase sims ("mean", "max", "topk_mean").

    Returns:
        {
          "ok": bool,
          "measured": {"bloom_best": str, "bloom_scores": dict, "dok_best": str, "dok_scores": dict},
          "feedback": str
        }
    """
    if _BACKEND is None or _BLOOM_INDEX is None or _DOK_INDEX is None:
        return {"error": "Classifier not initialized. Call set_classifier_state(backend, bloom_index, dok_index) first."}

    try:
        res = classify_levels_phrases(
            question,
            BLOOMS_PHRASES,
            DOK_PHRASES,
            backend=_BACKEND,
            prebuilt_bloom_index=_BLOOM_INDEX,
            prebuilt_dok_index=_DOK_INDEX,
            agg=agg,
            return_phrase_matches=True
        )
    except Exception as e:
        return {"error": f"classify_levels_phrases failed: {e}"}

    def _parse_target_bloom(t: str):
        order = ["Remember","Understand","Apply","Analyze","Evaluate","Create"]
        if t.endswith("+"):
            base = t[:-1]
            return set(order[order.index(base):])
        return {t}

    def _parse_target_dok(t: str):
        order = ["DOK1","DOK2","DOK3","DOK4"]
        if "-" in t:
            lo, hi = t.split("-")
            return set(order[order.index(lo):order.index(hi)+1])
        return {t}

    bloom_target_set = _parse_target_bloom(target_bloom)
    dok_target_set = _parse_target_dok(target_dok)

    bloom_best = res["blooms"]["best_level"]
    dok_best = res["dok"]["best_level"]

    bloom_ok = bloom_best in bloom_target_set
    dok_ok = dok_best in dok_target_set

    feedback_parts = []
    if not bloom_ok:
        feedback_parts.append(
            f"Shift Bloom’s from {bloom_best} toward {sorted(bloom_target_set)}. "
            f"Top cues: {res['blooms']['top_phrases'].get(bloom_best, [])[:3]}"
        )
    if not dok_ok:
        feedback_parts.append(
            f"Shift DOK from {dok_best} toward {sorted(dok_target_set)}. "
            f"Top cues: {res['dok']['top_phrases'].get(dok_best, [])[:3]}"
        )

    return {
        "ok": bool(bloom_ok and dok_ok),
        "measured": {
            "bloom_best": bloom_best,
            "bloom_scores": res["blooms"]["scores"],
            "dok_best": dok_best,
            "dok_scores": res["dok"]["scores"],
        },
        "feedback": " ".join(feedback_parts) if feedback_parts else "On target.",
    }