modules/cypher_cot_chain.py

"""CYPHER V12 M16 — Chain-of-thought multi-step reasoning.

For complex CYBERSEC prompts (CVE impact analysis, attack chain reconstruction,
multi-vuln correlation), decompose into atomic reasoning steps and chain them
through multiple generate calls.

Pipeline:
  1. Detect complex prompt (heuristic + keywords)
  2. Decompose into 2-4 sub-questions
  3. Generate answer for each sub-question (using same model)
  4. Assemble final structured response

The model itself doesn't have a "think" token (vocab 16384 small), so we
implement reasoning via multiple constrained calls.
"""
from __future__ import annotations

import logging
import re
from typing import Callable, Any

logger = logging.getLogger(__name__)


# Complex prompt indicators
_COMPLEX_KEYWORDS_CYBERSEC = (
    "kill chain", "attack chain", "exploit chain",
    "incident reconstruction", "post-mortem",
    "full analysis", "complete analysis",
    "step by step", "step-by-step",
    "impact and mitigation",
    "correlate", "correlation",
    "what to do if", "how would you respond",
    "incident response plan",
    "playbook",
    "what should I do about",
    "comment réagir si",
    "plan d'action",
    "analyser complètement",
)
_COMPLEX_KEYWORDS_TRADING = (
    "full setup analysis",
    "multi-timeframe analysis",
    "complete setup",
    "trade plan",
    "should i enter",
    "should i exit",
    "expected scenarios",
    "scenario analysis",
    "analyse complète du setup",
    "plan de trade",
)


def is_complex_prompt(prompt: str, category: str | None = None) -> bool:
    """Heuristic detection: is this prompt worth decomposing?"""
    if not prompt or len(prompt) < 20:
        return False
    p = prompt.lower()
    # CYBERSEC
    if category == "CYBERSEC" or category is None:
        if any(kw in p for kw in _COMPLEX_KEYWORDS_CYBERSEC):
            return True
    # TRADING
    if category == "TRADING" or category is None:
        if any(kw in p for kw in _COMPLEX_KEYWORDS_TRADING):
            return True
    # CVE + multiple concepts
    cve_count = len(re.findall(r"CVE-\d{4}-\d+", prompt, re.IGNORECASE))
    if cve_count >= 2:
        return True
    # Multiple "and/et" linking clauses
    if p.count(" and ") + p.count(" et ") >= 2 and len(p) > 100:
        return True
    return False


def decompose_cybersec(prompt: str) -> list[str]:
    """Decompose a complex cybersec prompt into atomic sub-questions."""
    cves = re.findall(r"CVE-\d{4}-\d+", prompt, re.IGNORECASE)
    techniques = re.findall(r"\bT\d{4}(?:\.\d{3})?\b", prompt, re.IGNORECASE)

    sub_questions: list[str] = []
    if cves:
        # Per-CVE analysis
        for cve in cves[:3]:
            sub_questions.append(f"What is {cve} and its severity?")
        if len(cves) >= 2:
            sub_questions.append(f"How do {cves[0]} and {cves[1]} relate or chain?")
        sub_questions.append("What are the mitigations for the discussed CVEs?")
    elif techniques:
        for tid in techniques[:3]:
            sub_questions.append(f"What is MITRE {tid} and how is it executed?")
        sub_questions.append("How to detect these MITRE techniques?")
    else:
        # Generic decomposition
        sub_questions = [
            "What is the threat being described?",
            "What are the technical mechanisms involved?",
            "How can it be detected in logs or telemetry?",
            "What are the recommended mitigations?",
        ]
    return sub_questions


def decompose_trading(prompt: str) -> list[str]:
    """Decompose a complex trading prompt."""
    return [
        "What is the higher-timeframe (HTF) market structure context?",
        "What is the key SMC pattern visible (Order Block, FVG, sweep)?",
        "What is the entry zone and confluence?",
        "Where should the stop loss and take profit be placed?",
    ]


def assemble_response(
    prompt: str,
    sub_questions: list[str],
    sub_answers: list[str],
    category: str | None = None,
) -> str:
    """Combine sub-answers into a structured final response."""
    parts: list[str] = []
    parts.append(f"Multi-step analysis of: {prompt[:120]}")
    parts.append("")
    for i, (q, a) in enumerate(zip(sub_questions, sub_answers), 1):
        parts.append(f"Step {i}: {q}")
        parts.append(f"  {a}")
        parts.append("")
    return "\n".join(parts)


class CoTChain:
    """Chain-of-thought wrapper around a model generate function."""

    def __init__(
        self,
        generate_fn: Callable[[str, int, float], str],
        max_sub_questions: int = 4,
        sub_max_tokens: int = 150,
        sub_temperature: float = 0.30,
    ):
        self.generate_fn = generate_fn
        self.max_sub_questions = max_sub_questions
        self.sub_max_tokens = sub_max_tokens
        self.sub_temperature = sub_temperature

    def run(self, prompt: str, category: str | None = None) -> dict:
        if not is_complex_prompt(prompt, category=category):
            # Direct single-shot
            single = self.generate_fn(prompt, 200, 0.35)
            return {
                "mode": "single",
                "response": single,
                "n_steps": 1,
            }

        # Decompose
        if category == "TRADING":
            sub_questions = decompose_trading(prompt)
        else:
            sub_questions = decompose_cybersec(prompt)
        sub_questions = sub_questions[: self.max_sub_questions]

        # Generate sub-answers
        sub_answers: list[str] = []
        for sq in sub_questions:
            try:
                ans = self.generate_fn(sq, self.sub_max_tokens, self.sub_temperature)
            except Exception as e:
                logger.warning(f"sub-question gen failed: {e}")
                ans = f"[Sub-step failed: {type(e).__name__}]"
            sub_answers.append(ans)

        # Assemble
        final = assemble_response(prompt, sub_questions, sub_answers, category=category)
        return {
            "mode": "chain",
            "response": final,
            "n_steps": len(sub_questions),
            "sub_questions": sub_questions,
            "sub_answers": sub_answers,
        }


__all__ = [
    "is_complex_prompt",
    "decompose_cybersec",
    "decompose_trading",
    "assemble_response",
    "CoTChain",
]


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("=== M16 cypher_cot_chain SMOKE ===")

    # Detection
    tests = [
        ("What is CVE-2021-44228?", "CYBERSEC", False),
        ("Full kill chain analysis of CVE-2021-44228 and CVE-2024-3400 with mitigations", "CYBERSEC", True),
        ("Step by step how to detect lateral movement T1021", "CYBERSEC", True),
        ("Complete setup analysis for BTCUSDT 4H with entry plan", "TRADING", True),
        ("Hello", "CONV", False),
    ]
    for p, c, expected in tests:
        got = is_complex_prompt(p, category=c)
        mark = "✓" if got == expected else "✗"
        print(f"  {mark} '{p[:50]}' ({c}) complex={got}")

    # Decompose
    print(f"\nDecompose cybersec multi-CVE:")
    sqs = decompose_cybersec("Full analysis of CVE-2021-44228 and CVE-2024-3400 with mitigations")
    for i, s in enumerate(sqs, 1):
        print(f"  {i}. {s}")

    print(f"\nDecompose trading:")
    sqs = decompose_trading("Full setup analysis for BTCUSDT 4H")
    for i, s in enumerate(sqs, 1):
        print(f"  {i}. {s}")

    # CoTChain with mock model
    canned_answers = [
        "Log4Shell CVE-2021-44228 RCE in Apache Log4j2 CVSS 10.0.",
        "CVE-2024-3400 Palo Alto PAN-OS unauthenticated RCE.",
        "Both chained: gain initial foothold via Log4Shell, pivot via Palo Alto.",
        "Mitigations: patch Log4j 2.17.1+ and PAN-OS hotfix.",
    ]
    idx = [0]
    def mock_gen(p, mt, t):
        a = canned_answers[idx[0] % len(canned_answers)]
        idx[0] += 1
        return a

    chain = CoTChain(generate_fn=mock_gen)
    result = chain.run("Full analysis of CVE-2021-44228 and CVE-2024-3400 with mitigations", category="CYBERSEC")
    print(f"\nCoT mode: {result['mode']}, steps: {result['n_steps']}")
    print(f"Response (first 400 chars):\n{result['response'][:400]}")

    # Single-shot fallback
    result2 = chain.run("Hello", category="CONV")
    print(f"\nSingle mode: {result2['mode']}, response: {result2['response'][:80]}")

    print("=== SMOKE PASS ===")