import json
import logging

logger = logging.getLogger(__name__)


class VLMCompilerEngine:
    """
    The Self-Compiling Heuristics Engine
    This engine leverages a massive VLM to analyze failures in the Zero-Latency Engine.
    It takes a screenshot + XML dump, finds the missing intent, and generates a new,
    blazing-fast deterministic Regex/XPath rule to be cached and executed next time.
    """

    def __init__(self, device):
        self.device = device

    def generate_heuristic(self, intent_description: str, context_xml: str) -> dict:
        """
        Calls the VLM to visually find the intent in the screen, then cross-reference it
        with the provided XML to generate a deterministic extraction rule.
        """
        # Sanitize intent to avoid confusing the LLM with python list syntax
        clean_intent = intent_description
        if "['" in clean_intent:
            clean_intent = clean_intent.replace("['", "").replace("']", "").replace("', '", " AND ")

        logger.warning(
            f"🧠 [Compiler Engine] Deterministic heuristic failed for: '{clean_intent}'. Synthesizing new rule...",
            extra={"color": "\x1b[1m\x1b[35m"},
        )

        from GramAddict.core.config import Config

        cfg = Config()
        if not hasattr(cfg, "args"):
            raise RuntimeError("Config().args not initialized — cannot resolve AI model. Fail Fast.")
        model = getattr(cfg.args, "ai_telepathic_model")
        url = getattr(cfg.args, "ai_telepathic_url")
        use_local = "11434" in url or "localhost" in url

        simplified_xml = self._simplify_xml(context_xml)

        # --- Model Trust Logging ---
        from GramAddict.core.benchmark_guard import BENCHMARKS_FILE

        trust_log = f"Using {model}"
        try:
            if os.path.exists(BENCHMARKS_FILE):
                with open(BENCHMARKS_FILE, "r") as f:
                    bench_data = json.load(f).get("models", {}).get(model, {})
                    if bench_data:
                        score = bench_data.get("telepathic_score", 0)
                        passed = "PASS" if bench_data.get("passed_all", False) else "FAIL"
                        unsuitable = bench_data.get("is_unsuitable", False)
                        trust_level = (
                            "HIGH"
                            if score >= 80 and not unsuitable
                            else "MEDIUM"
                            if score >= 50 and not unsuitable
                            else "LOW/UNSAFE"
                        )
                        trust_log += f" [Benchmark: {score}/100 | {passed} | Trust: {trust_level}]"
                        if unsuitable:
                            logger.error(f"⛔ [Safety Alert] {model} is marked as UNSUITABLE for this task!")
        except Exception:
            pass
        logger.info(f"🧠 [Compiler] Intent: '{clean_intent}' -> {trust_log}")
        # ---------------------------

        system_prompt = (
            "You are a regex engineering expert for Android UI. Identify the most stable resource-id regex for the intent.\n"
            "Rules:\n"
            "1. Output ONLY a raw JSON object.\n"
            "2. NO markdown, NO triple backticks.\n"
            '3. Format: {"rule_type": "regex", "target_attribute": "resource-id", "pattern": ".*regex.*", "confidence": 0.95, "reasoning": "string"}'
        )

        user_prompt = f"TARGET INTENT: {clean_intent}\n\nUI XML:\n{simplified_xml[:2000]}"

        try:
            from GramAddict.core.llm_provider import query_telepathic_llm

            res_text = query_telepathic_llm(
                model=model,
                url=url,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                temperature=0.1,
                use_local_edge=use_local,
            )

            if not res_text:
                logger.error("Compiler LLM returned empty response.")
                return None

            if "```json" in res_text:
                res_text = res_text.split("```json")[1].split("```")[0].strip()
            elif res_text.startswith("```"):
                res_text = "\n".join(res_text.strip().split("\n")[1:-1])

            try:
                decision = json.loads(res_text)
            except json.JSONDecodeError:
                logger.error(f"Compiler LLM returned invalid JSON: {res_text[:100]}...")
                return None

            # If LLM returned a list, take the first item if it's a dict
            if isinstance(decision, list):
                if len(decision) > 0 and isinstance(decision[0], dict):
                    decision = decision[0]
                else:
                    logger.error(f"Compiler LLM returned unexpected list format: {decision}")
                    return None

            if not isinstance(decision, dict):
                logger.error(f"Compiler LLM returned non-object response: {type(decision)}")
                return None

            pattern = decision.get("pattern")
            if not pattern:
                logger.error("Compiler LLM returned empty rule pattern. Aborting heuristic generation.")
                return None

            logger.info(
                f"✨ [Compiler] New Heuristic Synthesized! Rule: {decision.get('rule_type')} -> {pattern}",
                extra={"color": "\x1b[1m\x1b[32m"},
            )

            if decision.get("rule_type") == "xpath":
                logger.error(
                    "Compiler LLM returned 'xpath'. Rejecting rule because it causes xml.etree crashes. Will fallback/retry."
                )
                return None

            return {
                "rule_type": "regex",
                "target_attribute": decision.get("target_attribute", "text"),
                "pattern": pattern,
            }

        except Exception as e:
            logger.error(f"Heuristic compilation crashed: {e}")
            return None

    def _simplify_xml(self, xml_tree: str) -> str:
        import xml.etree.ElementTree as ET

        nodes = []
        try:
            root = ET.fromstring(xml_tree)
            for i, node in enumerate(root.iter("node")):
                attrib = node.attrib
                text = attrib.get("text", "")
                desc = attrib.get("content-desc", "")
                res_id = attrib.get("resource-id", "")
                if text or desc or res_id:
                    nodes.append(f"[{i}] text='{text}' desc='{desc}' r_id='{res_id}'")
        except:
            pass
        return "\n".join(nodes)
