import logging
from typing import List, Optional

from GramAddict.core.config import Config
from GramAddict.core.llm_provider import query_llm

logger = logging.getLogger(__name__)


def ask_brain_for_action(
    goal: str, screen_type: str, available_actions: List[str], explored_actions: set, context: dict = None
) -> Optional[str]:
    """Asks the VLM to decide the best available action to reach the goal, considering failures."""
    if not available_actions:
        return None

    cfg = Config()
    if not hasattr(cfg, "args"):
        raise RuntimeError("Config().args not initialized — cannot resolve AI model. Fail Fast.")
    url = getattr(cfg.args, "ai_model_url")
    model = getattr(cfg.args, "ai_model")

    prompt = (
        f"You are an autonomous Instagram agent. Your ultimate goal is: '{goal}'.\n"
        f"You are currently on the screen: {screen_type}.\n"
        f"These actions are available to you right now: {available_actions}\n"
    )
    if explored_actions:
        prompt += f"You recently tried these actions but they failed or didn't help: {list(explored_actions)}\n"
    if context:
        prompt += f"Context: {context}\n"

    prompt += (
        "INSTRUCTIONS:\n"
        "1. Reason about where you are. Consider the screen type and what actions make sense on that screen.\n"
        "2. If the goal requires navigating away from the current screen, choose the action that moves you closest to the goal.\n"
        "3. 'scroll down' reveals more UI elements on scrollable screens (feeds, profiles, lists). If your target is likely on this screen but not currently visible, you MUST choose 'scroll down'.\n"
        "4. 'press back' exits the current screen and returns to the previous one. Use it when you are on a screen that doesn't lead to your goal.\n"
        "5. DO NOT hallucinate actions. Reply ONLY with the exact string from the available actions list.\n"
        "6. Reply with ONLY the action string, nothing else."
    )

    try:
        response = query_llm(
            url=url,
            model=model,
            prompt="Choose the next best action.",
            system=prompt,
            format_json=False,
            max_tokens=250,
        )
        if response:
            result = response if isinstance(response, str) else response.get("response", "")
            result = result.strip().strip("'\"")

            # 1. Exact match check (ideal case)
            for act in available_actions:
                if act.lower() == result.lower():
                    return act
                    
            # 2. Strict line-by-line check (often the model outputs the action on the last line)
            for line in reversed(result.splitlines()):
                line = line.strip().strip("'\"")
                for act in available_actions:
                    if act.lower() == line.lower():
                        return act

            # 3. Fuzzy match (find the LAST mentioned action in the text, assuming it's the conclusion)
            best_act = None
            best_idx = -1
            for act in available_actions:
                idx = result.lower().rfind(act.lower())
                if idx > best_idx:
                    best_idx = idx
                    best_act = act
            
            if best_act:
                logger.warning(f"🧠 [Brain] Extracted action '{best_act}' from verbose LLM output.")
                return best_act

            logger.warning(f"🧠 [Brain] LLM returned an invalid action or no action found: '{result[:100]}...'. Falling back.")
    except Exception as e:
        logger.debug(f"🧠 [Brain] Error querying LLM: {e}")

    return None
