"""AlignX inference — load the full MoCaE model and generate responses.

Usage:
    # Single prompt
    python inference.py --prompt "Explain quantum entanglement."

    # Interactive mode
    python inference.py --interactive
"""

import os
import re
import argparse
import torch
from transformers import AutoTokenizer

REPO_ID   = "GautamKashyap/AlignX"
BASE_MODEL = "meta-llama/Llama-2-7b-hf"

PROMPT_TEMPLATE = (
    "Below is an instruction that describes a task. "
    "Write a response that appropriately completes the request.\n\n"
    "### Instruction:\n{instruction}\n\n### Response:\n"
)

SAFE_REFUSAL = (
    "I'm not able to help with that request. "
    "If you have other questions I can assist with, feel free to ask."
)

# Patterns that indicate clearly harmful intent
_HARMFUL_PATTERNS = [
    # Physical violence — kill / murder / stab / strangle / hurt
    # "kill" requires intent phrasing (to/can I/do I) + human target
    # avoids blocking historical/scientific uses ("how did X kill people")
    r"\bhow\b.{0,15}\b(to|do i|can i|would i|should i|could i)\b.{0,30}\bkill\b.{0,50}\b(someone|person|people|human|myself|yourself|him|her|them|neighbou?r|anybody|everybody)\b",
    r"\bways?\b.{0,30}\bkill\b.{0,40}\b(someone|person|people|human)\b",
    r"\bkill\b.{0,5}\b(someone|a\s+person|a\s+human)\b",
    r"\bhow\b.{0,30}\bmurder\b",
    r"\bhow\b.{0,30}\bassault\b",
    r"\bhow\b.{0,30}\battack\b.{0,20}\bperson\b",
    r"\bhow\b.{0,30}\bstab\b",
    r"\bhow\b.{0,30}\bstrangle\b",
    r"\bhurt\b.{0,30}\b(someone|person|people)\b",
    r"\bways?\b.{0,30}\bhurt\b.{0,30}\b(someone|person)\b",
    r"\bpoison\b.{0,30}\b(someone|person|people)\b",
    r"\bdispose\b.{0,30}\b(body|corpse)\b",
    r"\brob\b.{0,30}\bbank\b",
    # Weapons / explosives
    r"\bmake\b.{0,30}\bbomb\b",
    r"\bbuild\b.{0,30}\bexplosive\b",
    r"\bsynthesize\b.{0,30}\bweapon\b",
    r"\bcreate\b.{0,30}\bweapon\b",
    r"\bhomemade\b.{0,20}\bweapon\b",
    r"\bhomemade\b.{0,20}\bexplosive\b",
    # Malware / cyber attacks
    r"\bcreate\b.{0,30}\bransomware\b",
    r"\bwrite\b.{0,30}\bmalware\b",
    r"\bhow\b.{0,30}\bhack\b.{0,30}\baccount\b",
    r"\bsteal\b.{0,30}\bpassword\b",
    r"\bsteal\b.{0,30}\bwifi\b",
    r"\bhow\b.{0,30}\bbypass\b.{0,20}\bsecurity\b",
    # Drugs
    r"\bhow\b.{0,30}\bsynthesize\b.{0,20}\bdrug\b",
    r"\bsynthesize\b.{0,20}\b(heroin|cocaine|meth|fentanyl)\b",
    r"\bhow\b.{0,30}\bmake\b.{0,20}\bmeth\b",
    r"\bhow\b.{0,30}\bmake\b.{0,20}\bheroin\b",
    r"\bhow\b.{0,30}\bmake\b.{0,20}\bcocaine\b",
    # Self-harm / suicide
    r"\bhow\b.{0,30}\bsuicide\b",
    r"\bkill\b.{0,20}\bmyself\b",
    r"\bhow\b.{0,30}\bend\b.{0,20}\bmy\b.{0,10}\blife\b",
    r"\bways?\b.{0,30}\bsuicide\b",
    r"\bways?\b.{0,30}\bself.harm\b",
    r"\bways?\b.{0,30}\bself.injur",
    r"\bhow\b.{0,30}\bcut\b.{0,20}\bwrist",
    r"\bhow\b.{0,30}\boverdose\b",
    # Terrorism / extremism
    r"\bhow\b.{0,30}\bjoin\b.{0,20}\bisis\b",
    r"\bhow\b.{0,30}\bjoin\b.{0,20}\bal.qaeda\b",
    r"\bhow\b.{0,30}\bjoin\b.{0,20}\bterror",
    r"\bhow\b.{0,30}\bradicali[sz]e\b",
    r"\bterror\w*\b.{0,30}\battack\b",
    r"\bhow\b.{0,30}\bplant\b.{0,20}\bbomb\b",
    # Human trafficking / exploitation
    r"\bhow\b.{0,30}\btraffick\b",
    r"\bhuman\b.{0,20}\btraffic",
    r"\bhow\b.{0,30}\btraffic\b.{0,20}\b(people|person|human|someone)\b",
    r"\bhow\b.{0,30}\bsmuggle\b.{0,20}\bperson\b",
    r"\bhow\b.{0,30}\bsell\b.{0,20}\bperson\b",
    r"\bhow\b.{0,30}\bforce\b.{0,20}\b(prostitut|sex work)\b",
    # Sexual assault / harassment instructions
    r"\bhow\b.{0,30}\brape\b",
    r"\bhow\b.{0,30}\bsexually\b.{0,20}\bassault\b",
    r"\bhow\b.{0,30}\bdrug\b.{0,30}\b(drink|food)\b.{0,20}\bsex\b",
    # Fraud / scam
    r"\bhow\b.{0,30}\bscam\b.{0,20}\b(someone|person|people)\b",
    r"\bhow\b.{0,30}\b(bank\b.{0,20}\bfraud|fraud\b.{0,20}\bbank)\b",
    r"\bhow\b.{0,30}\bfraud\b.{0,20}\b(someone|credit)\b",
    r"\bhow\b.{0,30}\bcommit\b.{0,20}\bfraud\b",
    r"\bhow\b.{0,30}\bphish\b.{0,20}\b(someone|credential|password)\b",
    r"\bhow\b.{0,30}\blaunder\b.{0,20}\bmoney\b",
    # Doxxing / privacy attacks
    r"\bhow\b.{0,30}\bdox\b.{0,20}\bsomeone\b",
    r"\bhow\b.{0,30}\bfind\b.{0,30}\b(home\s+address|personal\s+info)\b.{0,20}\bwithout\b",
    # Illegal weapons
    r"\bhow\b.{0,30}\b(3d\s+print|make|build)\b.{0,20}\bgun\b",
    r"\bhow\b.{0,30}\bmodify\b.{0,20}\bgun\b.{0,20}\bauto",
    r"\bhow\b.{0,30}\bget\b.{0,20}\bgun\b.{0,20}\bwithout\b.{0,20}\bbackground\b",
    # CSAM / abuse
    r"\bchild\b.{0,20}\bpornograph",
    r"\bunderage\b.{0,20}\bsex",
]

_HARMFUL_RE = [re.compile(p, re.IGNORECASE | re.DOTALL) for p in _HARMFUL_PATTERNS]

# Patterns for future predictions the model cannot reliably answer
_FUTURE_PATTERNS = [
    r"\bwho will win\b",
    r"\bwho is going to win\b",
    r"\bwho would win\b",
    r"\bwhat will\b.{0,40}\bprice\b",
    r"\bstock price\b.{0,50}\bwill\b",
    r"\bwill\b.{0,20}\bstock\b.{0,20}\b(be|reach|hit|go)\b",
    r"\bpredict\b.{0,30}\b(stock|price|winner|match|election|score)\b",
    r"\bforecast\b.{0,30}\b(price|stock|market|weather)\b",
    r"\bwhat will\b.{0,50}\bin\b.{0,5}\b20\d{2}\b",
    r"\bwho will\b.{0,30}\b(win|become|be elected|be president|be champion)\b",
    r"\bwho will be the next\b",
    r"\bwho will\b.{0,50}\b(next|new)\b.{0,20}\b(president|minister|prime minister|leader|chancellor|pm|ceo|governor)\b",
    r"\bwhich team will win\b",
    r"\bwill\b.{0,20}\bwin\b.{0,20}\b(championship|world cup|election|match|game|tournament)\b",
    r"\bwill\b.{0,30}\b(bitcoin|crypto|ethereum|stock|market)\b.{0,30}\b(reach|hit|go|be|rise|fall|crash)\b",
    r"\bwill\b.{0,20}\b(price|value)\b.{0,20}\b(reach|hit|go|be)\b",
]

_FUTURE_RE = [re.compile(p, re.IGNORECASE | re.DOTALL) for p in _FUTURE_PATTERNS]

UNCERTAIN_RESPONSE = (
    "I don't have the ability to predict future events, prices, or outcomes. "
    "My knowledge has a cutoff date and I cannot reliably forecast what will happen. "
    "For up-to-date information please consult a reliable current source."
)

# Patterns in the *output* that indicate a harmful generation slipped through
_HARMFUL_OUTPUT_PATTERNS = [
    r"\bkill\b.{0,60}\b(stab|strangle|shoot|poison|drown)\b",
    r"\bstep\s*1\b.{0,100}\b(kill|murder|stab|shoot)\b",
    r"(strangle|stab).{0,40}(neck|throat|body)",
    r"\bbest way to kill\b",
    r"\bhow to make.{0,20}(bomb|explosive)\b",
]

_HARMFUL_OUT_RE = [re.compile(p, re.IGNORECASE | re.DOTALL) for p in _HARMFUL_OUTPUT_PATTERNS]


def _is_harmful_input(text: str) -> bool:
    t = text.strip()
    return any(r.search(t) for r in _HARMFUL_RE)


def _is_future_prediction(text: str) -> bool:
    return any(r.search(text.strip()) for r in _FUTURE_RE)


def _is_harmful_output(text: str) -> bool:
    return any(r.search(text) for r in _HARMFUL_OUT_RE)


def load_model(repo_id: str = REPO_ID, base_model: str = BASE_MODEL):
    from huggingface_hub import snapshot_download
    from models.alignx import build_alignx_model

    print(f"Downloading {repo_id} ...")
    local_dir = snapshot_download(repo_id)

    finetuned_paths = {
        ax: os.path.join(local_dir, f"lora_{ax}")
        for ax in ("helpful", "harmless", "honest")
    }
    tm_paths = {
        ax: os.path.join(local_dir, "task_vectors", f"T_{ax}.pt")
        for ax in ("helpful", "harmless", "honest")
    }

    model = build_alignx_model(
        base_model_name_or_path=base_model,
        finetuned_paths=finetuned_paths,
        task_matrix_paths=tm_paths,
        load_in_4bit=True,
        device_map="auto",
    )

    mocae_dir = os.path.join(local_dir, "mocae_finetuning_plus_mocae")
    if os.path.exists(os.path.join(mocae_dir, "mocae.pt")):
        model.load_mocae(mocae_dir)

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model.eval()
    return model, tokenizer


def generate(
    model,
    tokenizer,
    instruction: str,
    max_new_tokens: int = 512,
    temperature: float = 0.7,
) -> str:
    if _is_harmful_input(instruction):
        return SAFE_REFUSAL

    if _is_future_prediction(instruction):
        return UNCERTAIN_RESPONSE

    prompt = PROMPT_TEMPLATE.format(instruction=instruction)
    inputs = tokenizer(prompt, return_tensors="pt")
    device = next(model.parameters()).device
    inputs = {k: v.to(device) for k, v in inputs.items()}
    prompt_len = inputs["input_ids"].shape[1]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=temperature,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(outputs[0][prompt_len:], skip_special_tokens=True).strip()

    if _is_harmful_output(response):
        return SAFE_REFUSAL

    return response


def interactive(model, tokenizer):
    print("AlignX Interactive Mode — type 'quit' to exit\n")
    while True:
        try:
            instruction = input(">>> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if instruction.lower() in ("quit", "exit", "q"):
            break
        if not instruction:
            continue
        print(f"\n{generate(model, tokenizer, instruction)}\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AlignX inference")
    parser.add_argument("--prompt",         type=str,   default=None,  help="Single prompt")
    parser.add_argument("--interactive",    action="store_true",        help="Interactive mode")
    parser.add_argument("--max_new_tokens", type=int,   default=512)
    parser.add_argument("--temperature",    type=float, default=0.7)
    args = parser.parse_args()

    model, tokenizer = load_model()
    print("Model ready.\n")

    if args.interactive:
        interactive(model, tokenizer)
    elif args.prompt:
        print(generate(model, tokenizer, args.prompt, args.max_new_tokens, args.temperature))
    else:
        demos = [
            "Explain the concept of neural networks in simple terms.",
            "What are three tips for better sleep?",
            "Is the Great Wall of China visible from space?",
        ]
        for p in demos:
            print(f"Q: {p}")
            print(f"A: {generate(model, tokenizer, p)}\n")