#!/bin/bash
set -euo pipefail

FILES_LIST="${1:-text_files.txt}"
OUT_JSON="${2:-scan_results.json}"
OLLAMA_MODEL="${3:-llama3.1:8b}"
MAX_FILE_BYTES="${4:-200000}"
REQUEST_TIMEOUT_SEC="${5:-180}"

if [[ ! -s "$FILES_LIST" ]]; then
  echo "[]" > "$OUT_JSON"
  echo "[INFO] No text files to scan. Generated empty JSON: $OUT_JSON"
  exit 0
fi

TMP_DIR="$(mktemp -d)"
trap 'rm -rf "$TMP_DIR"' EXIT

SYS_PROMPT=$'You are a precise secure code scanning assistant. Analyze the provided file content and return findings as STRICT JSON only with this schema:\n{\n  "findings": [\n    {\n      "severity": "LOW|MEDIUM|HIGH|CRITICAL",\n      "rule_id": "string-identifier",\n      "line": number|null,\n      "description": "short human-readable description",\n      "recommendation": "concise fix guidance",\n      "references": ["optional references URLs or identifiers"]\n    }\n  ]\n}\nRules:\n- Output ONLY JSON. No markdown, no code fences, no commentary.\n- If no issues, output {\"findings\":[]}.\n- Prefer concrete, reproducible findings over theoretical. Highlight secrets, insecure configs, vulnerable patterns, risky dependencies (if visible), command injection, path traversal, deserialization, SSRF, XSS, SQLi, hardcoded credentials, bad crypto, weak TLS, etc.\n- Provide line numbers when confidently known; otherwise null.\n'

SYS_PROMPT=$(cat << EOF
You are a precise secure code scanning assistant. Analyze the provided file content and return findings as STRICT JSON only with this schema:
{
  "findings": [
    {
      "severity": "LOW|MEDIUM|HIGH|CRITICAL",
      "rule_id": "string-identifier",
      "line": number|null,
      "description": "short human-readable description",
      "recommendation": "concise fix guidance",
      "references": ["optional references URLs or identifiers"]
    }
  ]
}
Rules:
- Output ONLY JSON. No markdown, no code fences, no commentary.
- If no issues, output {\"findings\":[]}.
- Prefer concrete, reproducible findings over theoretical. Highlight secrets, insecure configs, vulnerable patterns, risky dependencies (if visible), command injection, path traversal, deserialization, SSRF, XSS, SQLi, hardcoded credentials, bad crypto, weak TLS, etc.
- Provide line numbers when confidently known; otherwise null.
EOF
)


export idx=0
for file in $(cat "$FILES_LIST"); do
  if [[ ! -f "$file" ]]; then
    echo "[WARN] Skipping missing file: $file"
    continue
  fi
  idx=$((idx+1))

  ext="${file##*.}"
  base="$(basename "$file")"
  code_tmp="$TMP_DIR/code_$idx.bin"
  head -c "$MAX_FILE_BYTES" "$file" | LC_ALL=C tr -d '\0' > "$code_tmp"
  bytes=$(wc -c < "$code_tmp" | tr -d ' ')

  req_json="$TMP_DIR/req_$idx.json"
  jq -n \
    --arg model "$OLLAMA_MODEL" \
    --arg sys "$SYS_PROMPT" \
    --arg f "$file" \
    --arg ext "$ext" \
    --rawfile code "$code_tmp" \
    '{
      model: $model,
      temperature: 0,
      stream: false,
      messages: [
        {role: "system", content: $sys},
        {role: "user", content: ("Analyze the following " +
                (if $ext == "" then "text" else $ext end) +
                " file named \"" + $f + "\" and return ONLY JSON as specified. File content between <<FILE>> markers:
<<FILE>>
" +
$code + "
<<FILE>>
")}
      ]
    }' > "$req_json"  
  
  echo "[INFO] [$idx] Scanning $file (${bytes} bytes) with model=$OLLAMA_MODEL ..."
  resp_json="$TMP_DIR/resp_$idx.json"
  set +e
  http_code=$(curl -sS -w "%{http_code}" -o "$resp_json" \
    --max-time "$REQUEST_TIMEOUT_SEC" --connect-timeout 10 \
    -H "Authorization: Bearer ${OLLAMA_API_KEY}" \
    -H "Content-Type: application/json" \
    -d @"$req_json" \
    "${OLLAMA_API_URI%/}/chat/completions")
  rc=$?
  set -e

  if [[ $rc -ne 0 || "$http_code" -lt 200 || "$http_code" -ge 300 ]]; then
    echo "[ERROR] Ollama request failed for $file (rc=$rc, http=$http_code)."
    # Create a synthetic "tool-error" entry
    echo "{\"file\":\"$file\",\"findings\":[{\"severity\":\"LOW\",\"rule_id\":\"tool-error\",\"line\":null,\"description\":\"Ollama request failed (rc=$rc, http=$http_code)\",\"recommendation\":\"Re-run scan or check Ollama endpoint/credentials\",\"references\":[]}],\"model\":\"$OLLAMA_MODEL\",\"bytes_analyzed\":$bytes}" \
      > "$TMP_DIR/file_$idx.json"
    continue
  fi
  
  # Extract assistant message content (model output)
  content=$(jq -r '.choices[0].message.content // ""' "$resp_json")

  # Try to parse content as JSON; if it has code fences, strip them
  parsed_ok=1
  echo "$content" | jq . > "$TMP_DIR/json_${idx}_raw.json" 2>/dev/null || parsed_ok=0
  if [[ $parsed_ok -eq 0 ]]; then
    content_stripped="$(printf "%s" "$content" | sed -E 's/^```(json)?[[:space:]]*//; s/```$//')"
    echo "$content_stripped" | jq . > "$TMP_DIR/json_${idx}_raw.json" 2>/dev/null || parsed_ok=0
  else
    parsed_ok=1
  fi

  if [[ $parsed_ok -eq 0 ]]; then
    echo "[WARN] Model did not return valid JSON for $file. Wrapping as tool-error."
    echo "{\"file\":\"$file\",\"findings\":[{\"severity\":\"LOW\",\"rule_id\":\"invalid-json\",\"line\":null,\"description\":\"Model returned non-JSON content\",\"recommendation\":\"Adjust prompt or model; ensure strict JSON output\",\"references\":[]}],\"model\":\"$OLLAMA_MODEL\",\"bytes_analyzed\":$bytes}" \
      > "$TMP_DIR/file_$idx.json"
  else
    # Normalize to object with findings array
    jq --arg file "$file" --arg model "$OLLAMA_MODEL" --argjson bytes "$bytes" '
      . as $orig
      | (if (has("findings") and (.findings|type=="array")) then $orig else {"findings": []} end)
      | . + {file:$file, model:$model, bytes_analyzed:$bytes}
    ' "$TMP_DIR/json_${idx}_raw.json" > "$TMP_DIR/file_$idx.json"
  fi
done

# Combine all per-file results into a single JSON array
jq -s '.' "$TMP_DIR"/file_*.json > "$OUT_JSON"
TOTAL=$(jq 'length' "$OUT_JSON")
echo "[INFO] Completed scan for ${TOTAL} file(s). Output: $OUT_JSON"
``