#!/usr/bin/env python3
import os, json
from pathlib import Path
from datetime import datetime
from collections import Counter

G="\033[92m"; Y="\033[93m"; R="\033[91m"; C="\033[96m"
W="\033[1m\033[97m"; DIM="\033[2m"; RST="\033[0m"; BLD="\033[1m"

def section(t):
    print(f"\n{C}{'─'*62}{RST}\n{W}  {t}{RST}\n{C}{'─'*62}{RST}")

def fmt_size(b):
    for u in ["B","KB","MB","GB"]:
        if b < 1024: return f"{b:.1f} {u}"
        b /= 1024
    return f"{b:.1f} TB"

def check_hf_cache():
    section("1 · HUGGINGFACE CACHED MODELS")
    cache = Path.home() / ".cache" / "huggingface" / "hub"
    if not cache.exists():
        print(f"  {Y}HF cache not found at {cache}{RST}"); return
    model_dirs = sorted(cache.glob("models--*"))
    if not model_dirs:
        print(f"  {Y}Cache empty{RST}"); return
    for md in model_dirs:
        name = md.name.replace("models--","").replace("--","/")
        total = sum(f.stat().st_size for f in md.rglob("*") if f.is_file())
        snaps = list((md/"snapshots").glob("*")) if (md/"snapshots").exists() else []
        snap_files = [f for s in snaps for f in (s.iterdir() if s.is_dir() else [])]
        has_config    = any("config.json"   in f.name for f in snap_files)
        has_tokenizer = any("tokenizer"     in f.name for f in snap_files)
        has_weights   = any(f.suffix in (".bin",".safetensors") for f in snap_files)
        print(f"\n  {W}{name}{RST}  {DIM}({fmt_size(total)}){RST}")
        print(f"    config.json : {G+'✓'+RST if has_config    else R+'✗ MISSING'+RST}")
        print(f"    tokenizer   : {G+'✓'+RST if has_tokenizer else R+'✗ MISSING'+RST}")
        print(f"    weights     : {G+'✓'+RST if has_weights   else R+'✗ MISSING'+RST}")
        print(f"    snapshots   : {G+str(len(snaps))+RST}")

def check_json_files():
    section("2 · JSON FILES")
    roots = [Path.home()/"leotsha_project", Path.home()]
    all_json = []
    for root in roots:
        if not root.exists(): continue
        for p in root.rglob("*.json"):
            if any(x in p.parts for x in [".cache","node_modules",".git","__pycache__"]): continue
            all_json.append(p)
    if not all_json:
        print(f"  {Y}No JSON files found{RST}"); return
    for p in sorted(all_json):
        st = p.stat()
        size = fmt_size(st.st_size)
        mod  = datetime.fromtimestamp(st.st_mtime).strftime("%Y-%m-%d %H:%M")
        print(f"\n  {W}{p.name}{RST}  {DIM}({size}, modified {mod}){RST}")
        print(f"  {DIM}{p}{RST}")
        try:
            with open(p, encoding="utf-8", errors="ignore") as f:
                data = json.load(f)
        except Exception as e:
            print(f"  {R}⚠  Parse error: {e}{RST}"); continue
        if isinstance(data, list):
            total = len(data)
            print(f"  Records : {BLD}{W}{total:,}{RST}")
            if total and isinstance(data[0], dict):
                print(f"  Keys    : {list(data[0].keys())}")
                lk = next((k for k in data[0] if k in ("label","sentiment","output")), None)
                if lk:
                    dist = Counter(str(r.get(lk,"?")) for r in data)
                    parts = [f"{G if 'pos' in l.lower() else R if 'neg' in l.lower() else Y}{l}{RST}:{c}" for l,c in dist.most_common()]
                    print(f"  Labels  : {' | '.join(parts)}")
                ok = G+"✓ ready"+RST if total>=500 else R+f"✗ need {500-total} more"+RST
                print(f"  Finetune ready (≥500): {ok}  {'  '+Y+'(2000+ recommended)'+RST if 500<=total<2000 else ''}")
        elif isinstance(data, dict):
            print(f"  Keys : {list(data.keys())[:8]}")

def check_local_models():
    section("3 · LOCAL MODEL CHECKPOINTS")
    found = []
    for root in [Path.home()/"leotsha_project", Path.home()/"models"]:
        if not root.exists(): continue
        for p in root.rglob("config.json"):
            if ".cache" not in str(p): found.append(p.parent)
    if not found:
        print(f"  {DIM}None found outside HF cache{RST}"); return
    for folder in found:
        files = list(folder.iterdir())
        names = {f.name for f in files}
        weights = [f for f in files if f.suffix in (".bin",".safetensors",".pt",".ckpt")]
        total = sum(f.stat().st_size for f in files if f.is_file())
        print(f"\n  {W}{folder}{RST}  ({fmt_size(total)})")
        print(f"    config.json      : {G+'✓'+RST if 'config.json' in names else R+'✗'+RST}")
        print(f"    tokenizer_config : {G+'✓'+RST if 'tokenizer_config.json' in names else R+'✗'+RST}")
        print(f"    weight files     : {G+str(len(weights))+' file(s)'+RST if weights else R+'none'+RST}")

print(f"\n{G}{'═'*62}{RST}")
print(f"{W}  Leotša la Sepedi — Resource & JSON Checker{RST}")
print(f"{G}{'═'*62}{RST}")
print(f"{DIM}  {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{RST}")
check_hf_cache()
check_json_files()
check_local_models()
print(f"\n{DIM}{'─'*62}{RST}\n")