import os
import json
import time
from typing import Dict, List, Any, Tuple

# Try importing Laya agent components
_LAYA_AGENT = None
_LAYA_AVAILABLE = False

try:
    from rl_agent_api import RLAgent
    from huggingface_hub import snapshot_download

    MODEL_REPO = os.environ.get("LAYA_MODEL_REPO", "convaiinnovations/laya")

    def _get_laya():
        global _LAYA_AGENT, _LAYA_AVAILABLE
        if _LAYA_AGENT is None:
            try:
                local_dir = snapshot_download(MODEL_REPO, token=os.environ.get("HF_TOKEN"))
                _LAYA_AGENT = RLAgent(local_dir, device="cpu")
                _LAYA_AVAILABLE = True
                print("✅ Laya System-1 Model successfully loaded on CPU.", flush=True)
            except Exception as e:
                print(f"⚠️ Laya direct model load deferred/fallback: {e}", flush=True)
                _LAYA_AVAILABLE = False
        return _LAYA_AGENT

except Exception as e:
    print(f"Laya modules import warning: {e}", flush=True)


class LayaDecisionOracle:
    """Laya System-1 Mathematical Decision & Gut-Feeling Oracle."""

    @staticmethod
    def decide_next_action(
        current_equations: List[str],
        goals: List[str],
        legal_actions: List[Dict[str, Any]],
        depth: int = 0,
    ) -> Tuple[Dict[str, Any], float, str]:
        """
        Uses Laya's System-1 typed choice decision to select the optimal mathematical action.
        Returns: (selected_action_dict, calibrated_confidence, instinct_reasoning)
        """
        if not legal_actions:
            return {"type": "finish", "description": "Finish"}, 1.0, "No other actions available."

        # If goal is already achieved, pick finish
        if any(eq.startswith("x =") or eq.startswith("x=") for eq in current_equations):
            for a in legal_actions:
                if a["type"] == "finish":
                    return a, 0.99, "Equations are isolated into solutions; declaring goal reached."

        # Formulate System 1 typed decision state
        state_repr = {
            "current_state": " and ".join(current_equations),
            "goal": " or ".join(goals) if goals else "solve for x",
            "search_depth": depth,
            "domain": "algebra",
        }

        criteria = {}
        action_map = {}
        for idx, act in enumerate(legal_actions):
            key = f"act_{idx}_{act['type']}"
            criteria[key] = act.get("description", act["type"])
            action_map[key] = act

        questions = {
            "math_action": {
                "type": "choice",
                "instructions": (
                    f"Given equation state `{state_repr['current_state']}` and goal `{state_repr['goal']}`, "
                    "which mathematical action will most rapidly reduce polynomial degree and isolate roots?"
                ),
                "criteria": criteria,
            }
        }

        # Try Laya single-forward-pass System 1 inference
        agent = None
        try:
            agent = _get_laya()
        except Exception:
            pass

        if agent is not None and _LAYA_AVAILABLE:
            try:
                t0 = time.perf_counter()
                results = agent.system_one(state_repr, questions)
                latency_ms = (time.perf_counter() - t0) * 1000

                choice_key = results["math_action"]["value"]
                confidence = float(results["math_action"].get("confidence", 0.85))
                selected = action_map.get(choice_key, legal_actions[0])

                reason = (
                    f"⚡ Laya System-1 forward pass ({latency_ms:.1f}ms): "
                    f"Selected [{selected['type']}] with {confidence*100:.1f}% calibrated confidence."
                )
                return selected, confidence, reason
            except Exception as e:
                print(f"Laya forward pass fallback: {e}", flush=True)

        # High-precision heuristic intuition fallback
        selected, conf, reason = LayaDecisionOracle._heuristic_gut_feeling(current_equations, legal_actions)
        return selected, conf, reason

    @staticmethod
    def _heuristic_gut_feeling(equations: List[str], legal_actions: List[Dict[str, Any]]) -> Tuple[Dict[str, Any], float, str]:
        """Mathematical instinct priority heuristic."""
        import re

        # 1. If all equations are already in isolated form, choose finish
        all_isolated = True
        for eq in equations:
            m = re.match(r'^\s*[a-zA-Z]\s*=\s*[+-]?\d+(?:\.\d+)?(?:\/\d+)?\s*$', eq)
            if not m:
                all_isolated = False
                break

        if all_isolated:
            for a in legal_actions:
                if a["type"] == "finish":
                    return a, 0.99, "⚡ Laya Instinct: All equation branches are isolated into canonical root forms; proving Q.E.D."

        # 2. If factored product exists (e.g. (x-2)*(x-3)=0), zero_product_split is highest priority
        for a in legal_actions:
            if a["type"] == "zero_product_split":
                return a, 0.97, "⚡ Laya Instinct: Factored product (A)(B)=0 detected; zero-product split decomposes problem into independent linear branches."

        # 3. If any equation can be factored, factor it
        for a in legal_actions:
            if a["type"] == "factor":
                return a, 0.94, "⚡ Laya Instinct: High-degree polynomial structure detected; factoring collapses nonlinear search space to linear components."

        # 4. If any un-isolated linear equation exists, isolate it
        for a in legal_actions:
            if a["type"] == "isolate":
                return a, 0.92, f"⚡ Laya Instinct: Linear term detected in {a.get('description', '')}; isolating directly yields root."

        # 5. If needs simplification
        for a in legal_actions:
            if a["type"] == "simplify":
                return a, 0.88, "⚡ Laya Instinct: Redundant algebraic terms detected; simplifying reduces complexity metric."

        return legal_actions[0], 0.80, f"⚡ Laya Instinct: Proposing {legal_actions[0]['type']} as optimal geodesic move."