"""
Inference script for Solidity Vulnerability Detector.

Loads the LoRA adapter from jhsu12/solidity-vulnerability-detector on top of
Qwen/Qwen2.5-Coder-7B-Instruct and generates a vulnerability assessment for
the given Solidity source code.

Usage:
    # Analyze a Solidity file:
    python inference.py --file contract.sol

    # Analyze inline code:
    python inference.py --code "pragma solidity ^0.8.0; contract Foo { ... }"

    # Interactive mode (paste code, press Ctrl-D to submit):
    python inference.py

    # Adjust generation parameters:
    python inference.py --file contract.sol --max_new_tokens 512 --temperature 0.1
"""

import argparse
import sys
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# ── Configuration ─────────────────────────────────────────────────────────────
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
ADAPTER_ID = "jhsu12/solidity-vulnerability-detector"

SYSTEM_PROMPT = (
    "You are a smart contract security auditor. Analyze the provided Solidity code "
    "for vulnerabilities.\n\n"
    "For each vulnerability found, report:\n"
    "- **Vulnerable**: Yes/No\n"
    "- **Type**: category (e.g. Reentrancy, Access Control, Integer Overflow/Underflow, "
    "Timestamp Dependence, Unchecked Low-Level Calls, tx.origin)\n"
    "- **Severity**: Critical/High/Medium/Low\n"
    "- **Location**: the affected function or line\n"
    "- **Impact**: what could go wrong\n"
    "- **Recommendation**: how to fix it\n\n"
    "If the contract is safe, state so clearly."
)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Run vulnerability detection inference on Solidity code."
    )
    parser.add_argument(
        "--file", type=str, default=None,
        help="Path to a .sol file to analyze"
    )
    parser.add_argument(
        "--code", type=str, default=None,
        help="Inline Solidity code string to analyze"
    )
    parser.add_argument(
        "--max_new_tokens", type=int, default=512,
        help="Maximum tokens to generate (default: 512)"
    )
    parser.add_argument(
        "--temperature", type=float, default=0.1,
        help="Sampling temperature (0 = greedy, default: 0.1)"
    )
    parser.add_argument(
        "--top_p", type=float, default=0.9,
        help="Top-p sampling (default: 0.9)"
    )
    parser.add_argument(
        "--load_in_4bit", action="store_true", default=True,
        help="Use 4-bit quantization (default: True, saves VRAM)"
    )
    parser.add_argument(
        "--load_in_8bit", action="store_true", default=False,
        help="Use 8-bit quantization instead of 4-bit"
    )
    return parser.parse_args()


def load_model(load_in_4bit=True, load_in_8bit=False):
    """Load the base model + LoRA adapter."""
    print(f"🤖 Loading base model: {BASE_MODEL}")
    print(f"🔌 LoRA adapter: {ADAPTER_ID}")

    # Device / dtype detection
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        has_bf16 = torch.cuda.is_bf16_supported()
        print(f"🖥️  GPU: {gpu_name} ({gpu_mem:.1f} GB)")
    else:
        has_bf16 = False
        print("⚠️  No GPU detected — running on CPU (will be slow)")

    compute_dtype = torch.bfloat16 if has_bf16 else torch.float16

    # Quantization config
    if load_in_8bit:
        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
        print("📦 Using 8-bit quantization")
    elif load_in_4bit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_use_double_quant=True,
        )
        print("📦 Using 4-bit quantization (NF4 + double quant)")
    else:
        bnb_config = None
        print("📦 No quantization (full precision)")

    # Attention implementation
    attn_impl = "sdpa"
    try:
        import flash_attn
        attn_impl = "flash_attention_2"
        print(f"⚡ Using flash_attention_2 (v{flash_attn.__version__})")
    except ImportError:
        print(f"⚡ Using SDPA (PyTorch native)")

    # Load base model
    model_kwargs = dict(
        device_map="auto",
        torch_dtype=compute_dtype,
        trust_remote_code=True,
        attn_implementation=attn_impl,
    )
    if bnb_config is not None:
        model_kwargs["quantization_config"] = bnb_config

    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, **model_kwargs)

    # Load LoRA adapter
    model = PeftModel.from_pretrained(model, ADAPTER_ID)
    model.eval()

    # Load tokenizer (from adapter repo — it may have custom chat template)
    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_ID, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    print("✅ Model loaded successfully\n")
    return model, tokenizer


def build_prompt(tokenizer, solidity_code):
    """Build a chat-formatted prompt for the model."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": solidity_code},
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    return text


def generate(model, tokenizer, solidity_code, max_new_tokens=512,
             temperature=0.1, top_p=0.9):
    """Generate a vulnerability assessment for the given Solidity code."""
    prompt = build_prompt(tokenizer, solidity_code)

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    input_len = inputs["input_ids"].shape[1]

    gen_kwargs = dict(
        max_new_tokens=max_new_tokens,
        pad_token_id=tokenizer.pad_token_id,
    )
    if temperature == 0 or temperature < 1e-6:
        gen_kwargs["do_sample"] = False
    else:
        gen_kwargs["do_sample"] = True
        gen_kwargs["temperature"] = temperature
        gen_kwargs["top_p"] = top_p

    with torch.no_grad():
        output_ids = model.generate(**inputs, **gen_kwargs)

    # Decode only the generated part (skip the prompt tokens)
    response = tokenizer.decode(output_ids[0][input_len:], skip_special_tokens=True)
    return response


def main():
    args = parse_args()

    # ── Get the Solidity code ─────────────────────────────────────────────────
    if args.file:
        print(f"📄 Reading Solidity file: {args.file}")
        with open(args.file, "r") as f:
            solidity_code = f.read()
    elif args.code:
        solidity_code = args.code
    else:
        print("📝 Paste your Solidity code below, then press Ctrl-D (Linux/Mac) "
              "or Ctrl-Z+Enter (Windows) to submit:\n")
        solidity_code = sys.stdin.read()

    if not solidity_code.strip():
        print("❌ No code provided. Use --file, --code, or pipe to stdin.")
        sys.exit(1)

    print(f"📏 Input code length: {len(solidity_code)} characters\n")

    # ── Load model ────────────────────────────────────────────────────────────
    model, tokenizer = load_model(
        load_in_4bit=args.load_in_4bit and not args.load_in_8bit,
        load_in_8bit=args.load_in_8bit,
    )

    # ── Generate ──────────────────────────────────────────────────────────────
    print("=" * 60)
    print("  VULNERABILITY ANALYSIS")
    print("=" * 60)

    response = generate(
        model, tokenizer, solidity_code,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
    )

    print(response)
    print("\n" + "=" * 60)


if __name__ == "__main__":
    main()