import argparse
import json
import os
import re
import sys
from datetime import datetime
from urllib.parse import quote_plus, urlparse, urlunparse
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError


REGISTRY_FILE = "model_registry.json"
URL_LOG_FILE = "model_url_scan.log"

USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/120.0.0.0 Safari/537.36"
)

CATEGORY_SPECS = [
    "checkpoints",
    "loras",
    "controlnets",
    "vae",
    "text_encoders",
    "diffusion_models",
    "clip",
    "embeddings",
]


def load_registry(path):
    if not os.path.exists(path):
        return None
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_registry(path, data):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def normalize_key(value):
    return re.sub(r"[^a-z0-9]+", "", (value or "").lower())


def filename_stem(filename):
    return os.path.splitext(filename or "")[0]


def is_huggingface_url(url):
    try:
        host = urlparse(url).netloc.lower()
    except Exception:
        return False
    return host.endswith("huggingface.co")


def is_civitai_url(url):
    try:
        host = urlparse(url).netloc.lower()
    except Exception:
        return False
    return host.endswith("civitai.com")


def normalize_huggingface_url(url, filename):
    parsed = urlparse(url.strip())
    scheme = parsed.scheme or "https"
    host = parsed.netloc
    path = parsed.path or ""

    path = path.replace("/blob/", "/resolve/")
    path = path.replace("/tree/", "/resolve/")

    clean_parts = [p for p in path.split("/") if p]
    if "resolve" in clean_parts:
        idx = clean_parts.index("resolve")
        repo_prefix = clean_parts[:idx]
        tail = clean_parts[idx + 1 :]

        branch = tail[0] if len(tail) >= 1 else "main"
        resolved_parts = repo_prefix + ["resolve", branch]
        if len(tail) >= 2:
            # keep any subpath, replace final filename with registry filename
            sub_parts = tail[1:-1] if len(tail) > 2 else []
            resolved_parts.extend(sub_parts)
        if filename:
            resolved_parts.append(filename)
        elif len(tail) >= 2:
            resolved_parts.append(tail[-1])
        else:
            resolved_parts.append("model.safetensors")
    else:
        if filename:
            resolved_parts = clean_parts + ["resolve", "main", filename]
        else:
            resolved_parts = clean_parts + ["resolve", "main"]

    new_path = "/" + "/".join(resolved_parts)
    return urlunparse((scheme, host, new_path, "", parsed.query, ""))


def extract_civitai_model_id(url):
    path = urlparse(url).path
    patterns = [
        r"/models/(\d+)",
        r"/api/v1/models/(\d+)",
    ]
    for pattern in patterns:
        m = re.search(pattern, path, flags=re.IGNORECASE)
        if m:
            return m.group(1)
    return None


def canonical_civitai_model_url(model_id):
    return f"https://civitai.com/models/{model_id}"


def fetch_json(url, timeout):
    req = Request(
        url,
        headers={
            "User-Agent": USER_AGENT,
            "Accept": "application/json",
        },
    )
    with urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8", errors="ignore")
    return json.loads(raw)


def extract_slug_from_url(url):
    parts = [p for p in urlparse(url).path.split("/") if p]
    if len(parts) >= 2 and parts[0].lower() == "models":
        token = parts[1]
        if token.isdigit():
            return ""
        return token
    return ""


def score_civitai_candidate(candidate, filename, stem_key, slug_key):
    score = 0
    c_name = str(candidate.get("name", ""))
    c_key = normalize_key(c_name)

    if c_key and (c_key in stem_key or stem_key in c_key):
        score += 20
    if slug_key and c_key and (slug_key in c_key or c_key in slug_key):
        score += 30

    versions = candidate.get("modelVersions") or []
    for version in versions:
        for file_info in (version.get("files") or []):
            f_name = str(file_info.get("name", ""))
            f_key = normalize_key(f_name)
            if not f_key:
                continue
            if f_name.lower() == filename.lower():
                score += 100
            elif stem_key and (stem_key in f_key or f_key in stem_key):
                score += 25

    return score


def resolve_civitai_model_url(current_url, filename, timeout, cache):
    existing_id = extract_civitai_model_id(current_url)
    if existing_id:
        return canonical_civitai_model_url(existing_id), "existing_id"

    stem = filename_stem(filename)
    stem_key = normalize_key(stem)
    slug = extract_slug_from_url(current_url)
    slug_key = normalize_key(slug)

    queries = []
    if slug:
        queries.append(slug)
    if stem:
        queries.append(stem)
    # unique while preserving order
    dedup = []
    seen = set()
    for q in queries:
        qn = q.strip()
        if not qn:
            continue
        low = qn.lower()
        if low in seen:
            continue
        seen.add(low)
        dedup.append(qn)

    best_url = None
    best_score = -1
    best_reason = ""

    for query in dedup:
        cache_key = query.lower()
        if cache_key in cache:
            data = cache[cache_key]
        else:
            api_url = f"https://civitai.com/api/v1/models?query={quote_plus(query)}&limit=30"
            try:
                data = fetch_json(api_url, timeout=timeout)
            except Exception:
                data = None
            cache[cache_key] = data

        if not isinstance(data, dict):
            continue

        for candidate in (data.get("items") or []):
            model_id = candidate.get("id")
            if not model_id:
                continue

            score = score_civitai_candidate(candidate, filename, stem_key, slug_key)
            if score > best_score:
                best_score = score
                best_url = canonical_civitai_model_url(model_id)
                best_reason = f"search:{query}"

    if best_url and best_score >= 30:
        return best_url, best_reason

    return current_url, "unresolved_slug"


def check_url_status(url, timeout):
    # HEAD first, then lightweight GET fallback if HEAD is blocked.
    headers = {"User-Agent": USER_AGENT}
    try:
        req = Request(url, headers=headers, method="HEAD")
        with urlopen(req, timeout=timeout) as resp:
            return int(resp.getcode() or 0), "HEAD"
    except HTTPError as err:
        return int(err.code), "HEAD"
    except Exception:
        pass

    try:
        headers["Range"] = "bytes=0-0"
        req = Request(url, headers=headers, method="GET")
        with urlopen(req, timeout=timeout) as resp:
            return int(resp.getcode() or 0), "GET_RANGE"
    except HTTPError as err:
        return int(err.code), "GET_RANGE"
    except URLError:
        return 0, "NETWORK_ERROR"
    except Exception:
        return 0, "ERROR"


def iter_registry_items(registry, checkpoints_only=False, name_filter=None):
    categories = ["checkpoints"] if checkpoints_only else CATEGORY_SPECS
    pattern = re.compile(name_filter, re.IGNORECASE) if name_filter else None
    for category in categories:
        items = registry.get(category, [])
        if not isinstance(items, list):
            continue
        for item in items:
            if isinstance(item, dict):
                if pattern:
                    name = str(item.get("name", ""))
                    if not pattern.search(name):
                        continue
                yield category, item


def sync_urls(
    registry_path,
    dry_run=False,
    offline=False,
    checkpoints_only=False,
    check_live=False,
    timeout=12,
    name_filter=None,
):
    registry = load_registry(registry_path)
    if not registry:
        print(f"Registry not found at: {registry_path}")
        return 1

    civitai_cache = {}
    scan_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_path = os.path.join(os.path.dirname(registry_path), URL_LOG_FILE)

    stats = {
        "scanned": 0,
        "changed": 0,
        "live_checked": 0,
        "live_ok": 0,
        "live_bad": 0,
        "unresolved": 0,
    }
    change_lines = []
    status_lines = []
    unresolved_lines = []

    for category, item in iter_registry_items(
        registry,
        checkpoints_only=checkpoints_only,
        name_filter=name_filter,
    ):
        name = str(item.get("name", "")).strip()
        old_url = str(item.get("url", "")).strip()
        if not name or not old_url:
            continue

        stats["scanned"] += 1
        new_url = old_url
        reason = "unchanged"

        if is_huggingface_url(old_url):
            new_url = normalize_huggingface_url(old_url, name)
            if new_url != old_url:
                reason = "huggingface_normalized"

        elif is_civitai_url(old_url):
            model_id = extract_civitai_model_id(old_url)
            if model_id:
                canonical = canonical_civitai_model_url(model_id)
                if canonical != old_url:
                    new_url = canonical
                    reason = "civitai_canonicalized"
            elif not offline:
                resolved, resolve_reason = resolve_civitai_model_url(
                    old_url,
                    name,
                    timeout=timeout,
                    cache=civitai_cache,
                )
                new_url = resolved
                reason = f"civitai_{resolve_reason}"
                if resolved == old_url and resolve_reason == "unresolved_slug":
                    stats["unresolved"] += 1
                    unresolved_lines.append(f"[{category}] {name}: {old_url}")

        else:
            reason = "unknown_provider"

        if new_url != old_url:
            stats["changed"] += 1
            item["url"] = new_url
            change_lines.append(f"[{category}] {name}: {old_url} -> {new_url} ({reason})")
            print(f"[URL FIXED] [{category}] {name}")

        if check_live:
            code, method = check_url_status(new_url, timeout=timeout)
            stats["live_checked"] += 1
            ok = 200 <= code < 400
            if ok:
                stats["live_ok"] += 1
            else:
                stats["live_bad"] += 1
            status_lines.append(f"[{category}] {name}: {code} via {method} | {new_url}")

    if dry_run:
        print("\nDry run complete. No changes written.")
    elif stats["changed"] > 0:
        save_registry(registry_path, registry)
        print(f"\nUpdated registry URLs: {stats['changed']} changes written.")
    else:
        print("\nNo URL changes needed.")

    with open(log_path, "w", encoding="utf-8") as log:
        log.write(f"Scan Date: {scan_time}\n")
        log.write(f"Registry: {registry_path}\n")
        log.write(f"Dry Run: {dry_run}\n")
        log.write(f"Offline: {offline}\n")
        log.write(f"Check Live: {check_live}\n")
        log.write(f"Checkpoints Only: {checkpoints_only}\n")
        log.write(f"Scanned: {stats['scanned']}\n")
        log.write(f"Changed: {stats['changed']}\n")
        log.write(f"Unresolved Slugs: {stats['unresolved']}\n")
        log.write(f"Live Checked: {stats['live_checked']}\n")
        log.write(f"Live OK: {stats['live_ok']}\n")
        log.write(f"Live Bad: {stats['live_bad']}\n")

        log.write("\n== URL Changes ==\n")
        for line in change_lines:
            log.write(f"{line}\n")

        log.write("\n== Live Status ==\n")
        for line in status_lines:
            log.write(f"{line}\n")

        log.write("\n== Unresolved Civitai Slugs ==\n")
        for line in unresolved_lines:
            log.write(f"{line}\n")

    print(f"URL scan log written to: {log_path}")
    print(
        "Summary: "
        f"scanned={stats['scanned']}, "
        f"changed={stats['changed']}, "
        f"unresolved={stats['unresolved']}, "
        f"live_checked={stats['live_checked']}, "
        f"live_bad={stats['live_bad']}"
    )

    return 0


def parse_args():
    parser = argparse.ArgumentParser(
        description="Validate and repair source download URLs in model_registry.json."
    )
    parser.add_argument("--registry", default=REGISTRY_FILE, help="Path to registry JSON.")
    parser.add_argument("--dry-run", action="store_true", help="Preview changes without writing.")
    parser.add_argument(
        "--offline",
        action="store_true",
        help="Skip network lookups (only normalize existing URL patterns).",
    )
    parser.add_argument(
        "--check-live",
        action="store_true",
        help="Check HTTP status for each final URL (HEAD + GET range fallback).",
    )
    parser.add_argument(
        "--checkpoints-only",
        action="store_true",
        help="Process checkpoints category only.",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=12,
        help="HTTP timeout in seconds for URL lookup/check requests.",
    )
    parser.add_argument(
        "--name-filter",
        default="",
        help="Optional regex filter for model names (example: nightvision).",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    registry_path = args.registry
    if not os.path.isabs(registry_path):
        registry_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), registry_path)

    rc = sync_urls(
        registry_path=registry_path,
        dry_run=args.dry_run,
        offline=args.offline,
        checkpoints_only=args.checkpoints_only,
        check_live=args.check_live,
        timeout=max(3, int(args.timeout)),
        name_filter=args.name_filter.strip() or None,
    )
    sys.exit(rc)