"""
Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.

This file is part of sunnypilot and is licensed under the MIT License.
See the LICENSE.md file in the root directory for more details.
"""

import sys
import hashlib
import json
import re
from pathlib import Path
from datetime import datetime, UTC


_VERSION_REGEX = re.compile(r'^(?:\d+[a-zA-Z]+|\d+[vVbB]\d+|[vVbB]\d+|\d{4})$')
_DECIMAL_VERSION_REGEX = re.compile(r'^([vVbB]?)(\d+(?:\.\d+)+)$')


def _version(word: str) -> str | None:
  match = _DECIMAL_VERSION_REGEX.match(word.strip('()[]{}.,;:'))
  if match:
    return match.group(1).upper() + match.group(2)
  stripped = re.sub(r'[^a-zA-Z0-9]', '', word)
  return stripped.upper() if _VERSION_REGEX.match(stripped) else None


def _truncate(short_name: str) -> str:
  return short_name[:8].rstrip('.')


def create_short_name(full_name: str) -> str:
  # Remove parentheses and extract alphanumeric words, keeping decimals for versioning
  clean_name = re.sub(r'\([^)]*\)', '', full_name)
  words = [re.sub(r'[^a-zA-Z0-9.]', '', word) for word in clean_name.split()]
  words = [word for word in words if word]

  if len(words) == 1:
    return _truncate((_version(words[0]) or words[0]).upper())

  # Handle special case: Name + Version (e.g., "Word A1" -> "WordA1", "Model v1.1" -> "ModelV1.1")
  if len(words) == 2 and re.match(r'^[A-Za-z]?\d+$', words[1]):
    return _truncate((words[0] + (_version(words[1]) or words[1])).upper())

  result = ""
  for word in words:
    # Version or number patterns
    version = _version(word)
    if version:
      result += version
    # All uppercase abbreviations (2-3 letters)
    elif re.match(r'^[A-Z]{2,3}$', word):
      result += word
    # Letters+digits (for example tr15 rev2)
    elif re.match(r'^[a-zA-Z]+[0-9]+$', word):
      result += word[0].upper() + ''.join(re.findall(r'\d+', word))
    elif word.isalpha():
      result += word[0].upper()
    elif word.isdigit():
      result += word
    else:
      result += word[0].upper()
  return _truncate(result)


def create_pkl_name(full_name: str) -> str:
  pkl = re.sub(r'[^a-zA-Z0-9]+', '_', full_name).strip('_').lower()
  return pkl


def _hash_pkl(pkl_path: Path) -> str:
  digest = hashlib.sha256()
  with pkl_path.open('rb') as f:
    while block := f.read(1024 * 1024):
      digest.update(block)
  return digest.hexdigest()


def _find_driving_pkl(output_path: Path) -> Path | None:
  for pattern in ('*driving_tinygrad.pkl', '*driving_*_tinygrad.pkl'):
    matches = sorted(output_path.glob(pattern))
    if matches:
      return matches[0]
  return None


def _hash_onnx_files(model_dir: Path) -> str | None:
  onnx_files = [f for f in sorted(model_dir.glob("*.onnx")) if f.stat().st_size > 1024]
  if not onnx_files:
    return None
  digest = hashlib.sha256()
  for f in onnx_files:
    with f.open('rb') as fh:
      while block := fh.read(1024 * 1024):
        digest.update(block)
  return digest.hexdigest()


def generate_model(driving_pkl: Path) -> dict:
  tinygrad_hash = _hash_pkl(driving_pkl)

  artifact_data = {
    "file_name": driving_pkl.name,
    "download_uri": {
      "url": "https://gitlab.com/sunnypilot/public/docs.sunnypilot.ai/-/raw/main/",
      "sha256": tinygrad_hash
    }
  }

  return {
    "type": "chunked",
    "artifact": artifact_data,
  }


def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown",
                         onnx_sha256=None, is_big=False) -> None:
  bundle_json = {
    "short_name": short_name,
    "display_name": custom_name or upstream_branch,
    "is_20hz": is_20hz,
    "ref": upstream_branch,
    "environment": "development",
    "runner": "tinygrad",
    "index": -1,
    "minimum_selector_version": "-1",
    "generation": "-1",
    "build_time": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
    "overrides": {},
    "is_big": is_big,
    "models": models,
  }

  if onnx_sha256:
    bundle_json["onnx_sha256"] = onnx_sha256

  # Write metadata to output_dir
  metadata_json = {
    "bundles": [bundle_json]
  }

  with open(output_dir / "metadata.json", "w") as f:
    json.dump(metadata_json, f, indent=2)
  print("Generated metadata.json")


if __name__ == "__main__":
  import argparse

  parser = argparse.ArgumentParser(description="Generate metadata JSON for the compiled JIT model")
  parser.add_argument("--model-dir", default="./models", help="Directory containing the model files")
  parser.add_argument("--output-dir", default="./output", help="Output directory for metadata")
  parser.add_argument("--custom-name", help="Custom display name for the model")
  parser.add_argument("--is-20hz", action="store_true", help="Whether this is a 20Hz model")
  parser.add_argument("--upstream-branch", default="unknown", help="Upstream branch name")
  parser.add_argument("--onnx-sha256", default=None, help="Pre-computed LFS SHA256 for the model file (used when no ONNX is present)")
  args = parser.parse_args()

  _output_dir = Path(args.output_dir)
  _output_dir.mkdir(exist_ok=True, parents=True)
  _short_name = create_short_name(args.custom_name) if args.custom_name else None
  _pkl = create_pkl_name(args.custom_name) if args.custom_name else None

  _driving_pkl = _find_driving_pkl(_output_dir)
  if not _driving_pkl:
    print(f"No driving_tinygrad.pkl found in {_output_dir}", file=sys.stderr)
    sys.exit(1)

  is_big = _driving_pkl.name.startswith('big_')

  if _pkl:
    new_pkl = _output_dir / f"driving_{_pkl}_tinygrad.pkl"
    if not new_pkl.exists():
      _driving_pkl = _driving_pkl.rename(new_pkl)
    else:
      _driving_pkl = new_pkl

  _model_metadata = generate_model(_driving_pkl)
  _onnx_sha256 = args.onnx_sha256 or _hash_onnx_files(Path(args.model_dir))
  create_metadata_json([_model_metadata], _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch,
                       onnx_sha256=_onnx_sha256, is_big=is_big)