"""CYPHER V12 M21 — Adaptive Compute Router.

Routes prompts to optimal inference strategy based on complexity:
  - EASY: single-shot greedy (1 sample)
  - MEDIUM: BoN=3 with M19
  - COMPLEX: BoN=8 + M16 cot_chain decomposition + M19 BoN per sub
  - VERY_COMPLEX: MCTS tree (M33 future)

Saves compute on easy prompts, spends compute on hard ones.

Complexity heuristics:
  - Length (>200 chars → +1)
  - Multiple CVEs/T-IDs → +1 per
  - "and/et" linkage → +0.5 each
  - Complex keywords (chain, analyze fully, step by step) → +2
  - Multiple categories detected → +1
  - Multi-turn context (PREV_C in prompt) → +0.5
"""
from __future__ import annotations

import logging
import re
from typing import Any

logger = logging.getLogger(__name__)


# Complexity indicators
_COMPLEX_KW = (
    "kill chain", "attack chain", "exploit chain", "step by step", "step-by-step",
    "complete analysis", "full analysis", "fully analyze", "fully analyse",
    "post-mortem", "post mortem", "incident reconstruction", "playbook",
    "scenario analysis", "multi-timeframe", "full setup",
    "analyse complète", "analyser complètement", "plan d'action",
    "in detail", "thoroughly", "exhaustively", "exhaustive",
)

_VERY_COMPLEX_KW = (
    "research", "investigate the full", "deep dive",
    "give me a complete report",
    "compare and contrast multiple",
)


class ComputeStrategy:
    """Inference strategy descriptor."""
    EASY = "easy"
    MEDIUM = "medium"
    COMPLEX = "complex"
    VERY_COMPLEX = "very_complex"


_STRATEGY_CONFIG = {
    ComputeStrategy.EASY: {
        "n_samples": 1,
        "temperature": 0.30,
        "top_k": 5,
        "use_cot_chain": False,
        "early_stop_score": 999.0,  # no early stop, single shot
        "max_new_tokens": 150,
    },
    ComputeStrategy.MEDIUM: {
        "n_samples": 3,
        "temperature": None,  # diverse profiles
        "top_k": None,
        "use_cot_chain": False,
        "early_stop_score": 4.5,
        "max_new_tokens": 200,
    },
    ComputeStrategy.COMPLEX: {
        "n_samples": 8,
        "temperature": None,
        "top_k": None,
        "use_cot_chain": True,
        "early_stop_score": 4.5,
        "max_new_tokens": 250,
    },
    ComputeStrategy.VERY_COMPLEX: {
        "n_samples": 8,
        "temperature": None,
        "top_k": None,
        "use_cot_chain": True,
        "use_mcts": True,  # placeholder for M33 future
        "early_stop_score": 4.7,
        "max_new_tokens": 350,
    },
}


def compute_complexity_score(prompt: str, category: str | None = None) -> tuple[float, dict]:
    """Numeric complexity 0..10 + breakdown."""
    if not prompt:
        return 0.0, {"empty": True}
    p_lower = prompt.lower()
    score = 0.0
    detail: dict = {}

    # Length
    if len(prompt) > 400:
        score += 2.0
        detail["very_long"] = len(prompt)
    elif len(prompt) > 200:
        score += 1.0
        detail["long"] = len(prompt)
    elif len(prompt) > 100:
        score += 0.5
        detail["medium_long"] = len(prompt)

    # CVEs
    n_cves = len(re.findall(r"cve-\d{4}-\d+", p_lower))
    if n_cves >= 2:
        score += 1.5 * n_cves
        detail["cves"] = n_cves
    elif n_cves == 1:
        score += 0.5
        detail["cves"] = 1

    # MITRE T-IDs
    n_tids = len(re.findall(r"\bt\d{4}", p_lower))
    if n_tids >= 2:
        score += 1.0 * n_tids
        detail["tids"] = n_tids
    elif n_tids == 1:
        score += 0.3
        detail["tids"] = 1

    # "and / et" linking
    n_and = p_lower.count(" and ") + p_lower.count(" et ")
    if n_and >= 2:
        score += 0.5 * n_and
        detail["and_links"] = n_and

    # Complex keywords
    cmplx_hits = sum(1 for k in _COMPLEX_KW if k in p_lower)
    if cmplx_hits > 0:
        score += 2.0 * cmplx_hits
        detail["complex_kw"] = cmplx_hits

    very_cmplx_hits = sum(1 for k in _VERY_COMPLEX_KW if k in p_lower)
    if very_cmplx_hits > 0:
        score += 3.0 * very_cmplx_hits
        detail["very_complex_kw"] = very_cmplx_hits

    # Multi-turn context
    if "[prev_u:" in p_lower or "[prev_c:" in p_lower or "[conv_hist]" in p_lower:
        score += 0.5
        detail["multi_turn"] = True

    # Category hint
    if category in ("CYBERSEC", "TRADING") and len(prompt) > 100:
        score += 0.5
        detail["domain_long"] = True

    return min(10.0, score), detail


def classify_strategy(prompt: str, category: str | None = None) -> tuple[str, float, dict]:
    """Return (strategy_name, complexity_score, detail)."""
    score, detail = compute_complexity_score(prompt, category)
    if score >= 6.0:
        strat = ComputeStrategy.VERY_COMPLEX
    elif score >= 3.0:
        strat = ComputeStrategy.COMPLEX
    elif score >= 1.0:
        strat = ComputeStrategy.MEDIUM
    else:
        strat = ComputeStrategy.EASY
    return strat, score, detail


def get_strategy_config(strategy: str) -> dict:
    return dict(_STRATEGY_CONFIG.get(strategy, _STRATEGY_CONFIG[ComputeStrategy.MEDIUM]))


class AdaptiveComputeRouter:
    """Orchestrates inference strategy selection + execution.

    Wires M19 (BoN), M22 (PRM), and optionally M16 (CoT decomposition).
    """

    def __init__(
        self,
        bon=None,            # TestTimeSamplingBoN instance
        prm=None,            # ProcessRewardModel instance
        cot_chain=None,      # CoTChain instance from M16
    ):
        self.bon = bon
        self.prm = prm
        self.cot_chain = cot_chain

    def run(
        self,
        prompt: str,
        category: str | None = None,
        force_strategy: str | None = None,
    ) -> dict:
        if force_strategy is not None:
            strategy = force_strategy
            cmplx_score = -1.0
            cmplx_detail = {"forced": True}
        else:
            strategy, cmplx_score, cmplx_detail = classify_strategy(prompt, category)
        cfg = get_strategy_config(strategy)
        result_meta = {
            "strategy": strategy,
            "complexity_score": cmplx_score,
            "complexity_detail": cmplx_detail,
            "config": cfg,
        }

        # Branch
        if strategy in (ComputeStrategy.COMPLEX, ComputeStrategy.VERY_COMPLEX) and self.cot_chain:
            if hasattr(self.cot_chain, "run"):
                cot_result = self.cot_chain.run(prompt, category=category)
                result_meta["cot_used"] = True
                result_meta["n_subqueries"] = cot_result.get("n_steps", 1)
                # If cot_chain decomposed (mode=chain), return its assembled response
                if cot_result.get("mode") == "chain":
                    return {
                        "response": cot_result["response"],
                        "meta": {**result_meta, "via": "cot_chain"},
                    }

        if self.bon is not None and cfg["n_samples"] > 1:
            bon_result = self.bon.run(
                prompt,
                category=category,
                n=cfg["n_samples"],
                max_new_tokens=cfg["max_new_tokens"],
            )
            return {
                "response": bon_result.best.response,
                "meta": {
                    **result_meta,
                    "via": "bon",
                    "n_sampled": bon_result.n_sampled,
                    "best_score": bon_result.best.score,
                    "all_scores": [c.score for c in bon_result.all_candidates],
                    "total_latency_ms": bon_result.total_latency_ms,
                },
            }

        # EASY or no BoN configured: single shot via generate_fn directly
        if self.bon is not None:
            single = self.bon.generate_fn(prompt, cfg["max_new_tokens"], cfg["temperature"] or 0.30,
                                          cfg.get("top_k") or 10, 0) if hasattr(self.bon, "generate_fn") else ""
        else:
            single = "[ERROR: no generate_fn available]"
        return {
            "response": single,
            "meta": {**result_meta, "via": "single_shot"},
        }


__all__ = [
    "ComputeStrategy",
    "compute_complexity_score",
    "classify_strategy",
    "get_strategy_config",
    "AdaptiveComputeRouter",
]


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("=== M21 cypher_adaptive_compute SMOKE ===")

    tests = [
        ("Hello", "CONV"),
        ("Who are you?", "IDENTITY"),
        ("What is CVE-2021-44228?", "CYBERSEC"),
        ("Full kill chain analysis of CVE-2021-44228 and CVE-2024-3400 with mitigations", "CYBERSEC"),
        ("Step-by-step incident reconstruction including lateral movement T1021, T1059, and exfiltration via DNS T1048.003", "CYBERSEC"),
        ("Complete setup analysis for BTCUSDT 4H with entry plan, multi-timeframe alignment, and risk management", "TRADING"),
        ("Research and deep dive on full APT29 TTPs with kill chain", "CYBERSEC"),
    ]
    for prompt, cat in tests:
        strat, score, detail = classify_strategy(prompt, cat)
        print(f"\n  '{prompt[:60]}'")
        print(f"  ↳ strategy={strat} score={score:.1f} detail={detail}")
        cfg = get_strategy_config(strat)
        print(f"  ↳ config: n={cfg['n_samples']} cot={cfg.get('use_cot_chain', False)} max_tok={cfg['max_new_tokens']}")

    # End-to-end with mocks
    from cypher_prm import ProcessRewardModel
    from cypher_tts_best_of_n import TestTimeSamplingBoN

    canned = ["I am CYPHER, the cybersec ASI of the family.",
              "CVE-2021-44228 (Log4Shell) is critical RCE Apache Log4j2 CVSS 10.0.",
              "Order Block is institutional candle imbalance, SMC entry."]
    idx = [0]
    def mock_gen(p, mt, t, tk=10, seed=0):
        out = canned[(idx[0] + seed) % len(canned)]
        idx[0] += 1
        return out

    prm = ProcessRewardModel()
    bon = TestTimeSamplingBoN(generate_fn=mock_gen,
                              scorer_fn=lambda r, p, c: prm.score(r, p, c),
                              default_n=3)
    router = AdaptiveComputeRouter(bon=bon, prm=prm, cot_chain=None)

    res = router.run("Who are you?", "IDENTITY")
    print(f"\nIDENTITY: strategy={res['meta']['strategy']} via={res['meta']['via']}")
    print(f"  Response: {res['response'][:80]}")

    res = router.run("Full kill chain analysis of CVE-2021-44228", "CYBERSEC")
    print(f"\nCYBERSEC complex: strategy={res['meta']['strategy']} via={res['meta']['via']}")

    print("\n=== SMOKE PASS ===")