import argparse
import json
import os
from pathlib import Path

import requests

# slack workflow setting can be edited at https://slack.com/shortcuts/Ft0AAER075L3/a314cf6f4a81543cc37b75702c5d10f1
SLACK_WEBHOOK_URL = os.environ.get("TRTLLM_PLC_WEBHOOK")
# this json file will be generated from pulse in pipeline scanning
INPUT_FILE = "./nspect_scan_report.json"

parser = argparse.ArgumentParser()
parser.add_argument("--build-url", required=True, help="Jenkins build URL")

args = parser.parse_args()

# Throw error if not set
if not SLACK_WEBHOOK_URL:
    raise EnvironmentError("Error: Environment variable 'TRTLLM_PLC_WEBHOOK' is not set!")

# Read file
raw_input = Path(INPUT_FILE).read_text()
vulnerabilities = json.loads(raw_input)


def safe(value, default="N/A"):
    return value if value else default


# Build attachment text
message_lines = ["* TensorRT LLM Source Code Vulnerability Scan Report*\n"]

severity_rank = {"Critical": 4, "High": 3, "Medium": 2, "Low": 1}

for v in vulnerabilities:
    sev = v.get("Severity", "Low")
    if severity_rank.get(sev, 0) <= 2:
        continue
    shortTermVersion = safe(v.get("Upgrade-Guidance", {}).get("Short-Term"))
    longTermVersion = safe(v.get("Upgrade-Guidance", {}).get("Long-Term"))
    lines = [
        f"🔴 *{safe(v.get('Severity'))}* — *{safe(v.get('Package Name'))}* `{safe(v.get('Package Version'))}`",
        f"• *CVE:* {safe(v.get('Related Vuln'))} | *BDSA:* {safe(v.get('CVE ID'))}",
        f"• *Score:* {safe(v.get('Score'))}",
        f"• *Status:* {safe(v.get('Status'))}",
        f"• *Published:* {safe(v.get('Vulnerability Published Date'))}",
        f"• *Upgrade:* `{shortTermVersion}` → `{longTermVersion}`",
        "─" * 40,  # separator line
    ]
    message_lines.extend(lines)

message_text = "\n".join(message_lines)

payload = {"report": message_text, "pipelineUrl": args.build_url}

print(payload)
# Send to Slack
resp = requests.post(SLACK_WEBHOOK_URL, json=payload, timeout=60)
resp.raise_for_status()