import gradio as gr
import requests
import json
import os
from datetime import datetime
import io
from huggingface_hub import HfApi

GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN", "")
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HF_ACCESS_TOKEN", "")

BUCKETS = [
    "GigaBites-AI-Team/gitback-cache-bucket-1",
    "MegaBites-AI/gitback-cache-bucket-2",
    "KiloBites-ai/gitback-cache-bucket-3",
]

def get_api():
    return HfApi(token=HF_TOKEN)

def github_headers():
    h = {"Accept": "application/vnd.github+json"}
    if GITHUB_TOKEN:
        h["Authorization"] = f"Bearer {GITHUB_TOKEN}"
    return h

def hf_headers():
    if HF_TOKEN:
        return {"Authorization": f"Bearer {HF_TOKEN}"}
    return {}

def cache_key(owner, repo, sha):
    return f"{owner}__{repo}/{sha}.json"

def index_key(owner, repo):
    return f"{owner}__{repo}/_index.json"

def check_cache(owner, repo, sha):
    key = cache_key(owner, repo, sha)
    for bucket in BUCKETS:
        url = f"https://huggingface.co/datasets/{bucket}/resolve/main/{key}?raw=true"
        r = requests.get(url, headers=hf_headers(), timeout=8)
        if r.status_code == 200:
            try:
                return r.json(), bucket
            except:
                continue
    return None, None

def get_index(owner, repo):
    key = index_key(owner, repo)
    primary_bucket = BUCKETS[hash(f"{owner}/{repo}") % len(BUCKETS)]
    url = f"https://huggingface.co/datasets/{primary_bucket}/resolve/main/{key}?raw=true"
    r = requests.get(url, headers=hf_headers(), timeout=8)
    if r.status_code == 200:
        try:
            return r.json(), primary_bucket
        except:
            pass
    return {}, primary_bucket

def update_index(owner, repo, sha, bucket):
    index, primary_bucket = get_index(owner, repo)
    if sha in index:
        return
    index[sha] = {"bucket": bucket, "cached_at": datetime.utcnow().isoformat()}
    key = index_key(owner, repo)
    try:
        api = get_api()
        content = json.dumps(index, indent=2).encode()
        api.upload_file(
            path_or_fileobj=io.BytesIO(content),
            path_in_repo=key,
            repo_id=primary_bucket,
            repo_type="dataset",
            commit_message=f"Index: {owner}/{repo}@{sha[:7]}"
        )
    except Exception as e:
        print(f"Index update failed: {e}")

def save_to_cache(owner, repo, sha, data):
    # Final dedup check
    existing, existing_bucket = check_cache(owner, repo, sha)
    if existing:
        print(f"Dedup: {sha[:7]} already in {existing_bucket}, skipping.")
        return existing_bucket

    # Deterministic bucket — same SHA always maps to same bucket
    bucket = BUCKETS[hash(sha) % len(BUCKETS)]
    key = cache_key(owner, repo, sha)

    try:
        api = get_api()
        content = json.dumps(data, indent=2).encode()
        api.upload_file(
            path_or_fileobj=io.BytesIO(content),
            path_in_repo=key,
            repo_id=bucket,
            repo_type="dataset",
            commit_message=f"Snapshot: {owner}/{repo}@{sha[:7]}"
        )
        print(f"✅ Cached {sha[:7]} → {bucket}")
        update_index(owner, repo, sha, bucket)
    except Exception as e:
        print(f"❌ Cache write failed: {e}")
        return f"error: {e}"

    return bucket

def fetch_commits(owner, repo):
    url = f"https://api.github.com/repos/{owner}/{repo}/commits?per_page=30"
    r = requests.get(url, headers=github_headers(), timeout=10)
    if r.status_code != 200:
        return None, f"❌ Error: {r.json().get('message', 'Unknown error')}"
    return r.json(), None

def fetch_snapshot(owner, repo, sha):
    # Check cache first
    cached, bucket = check_cache(owner, repo, sha)
    if cached:
        cached["_from_cache"] = True
        cached["_cache_bucket"] = bucket
        return cached

    # Fetch from GitHub
    r = requests.get(
        f"https://api.github.com/repos/{owner}/{repo}/commits/{sha}",
        headers=github_headers(), timeout=10
    )
    if r.status_code != 200:
        return None

    data = r.json()
    snapshot = {
        "sha": data["sha"],
        "message": data["commit"]["message"],
        "author": data["commit"]["author"]["name"],
        "date": data["commit"]["author"]["date"],
        "files": [
            {
                "filename": f["filename"],
                "status": f["status"],
                "additions": f.get("additions", 0),
                "deletions": f.get("deletions", 0),
                "patch": f.get("patch", "")[:2000]
            }
            for f in data.get("files", [])
        ],
        "_from_cache": False,
    }

    bucket = save_to_cache(owner, repo, sha, snapshot)
    snapshot["_cache_bucket"] = bucket
    return snapshot

def search_repo(repo_input):
    repo_input = repo_input.strip().strip("/")
    parts = repo_input.replace("https://github.com/", "").split("/")
    if len(parts) < 2:
        return gr.update(choices=[], value=None), "⚠️ Enter a valid repo like `owner/repo`", ""

    owner, repo = parts[0], parts[1]
    commits, err = fetch_commits(owner, repo)
    if err:
        return gr.update(choices=[], value=None), err, ""

    choices = [
        f"{c['sha'][:7]} — {c['commit']['message'][:60]} ({c['commit']['author']['name']})"
        for c in commits
    ]
    shas = [c["sha"] for c in commits]

    return (
        gr.update(choices=choices, value=None),
        f"✅ Found **{len(commits)}** recent commits for `{owner}/{repo}`",
        json.dumps({"owner": owner, "repo": repo, "shas": shas})
    )

def view_snapshot(selection, state_json):
    if not selection or not state_json:
        return "Select a commit above to view its snapshot."
    state = json.loads(state_json)
    idx = [s[:7] for s in state["shas"]].index(selection[:7])
    sha = state["shas"][idx]
    owner, repo = state["owner"], state["repo"]

    snap = fetch_snapshot(owner, repo, sha)
    if not snap:
        return "❌ Could not fetch snapshot."

    if snap.get("_from_cache"):
        source = f"⚡ **Served from cache** (`{snap.get('_cache_bucket', '')}`) — instant load!"
    else:
        source = f"🌐 **Fetched live from GitHub** → saved to `{snap.get('_cache_bucket', 'cache')}` (dedup guaranteed)"

    out = f"""## 📸 Snapshot: `{snap['sha'][:7]}`

{source}

**Message:** {snap['message']}
**Author:** {snap['author']}
**Date:** {snap['date']}
**Files changed:** {len(snap['files'])}

---
"""
    for f in snap["files"][:10]:
        out += f"\n### `{f['filename']}` — {f['status']} (+{f['additions']} / -{f['deletions']})\n"
        if f["patch"]:
            out += f"```diff\n{f['patch']}\n```\n"

    if len(snap["files"]) > 10:
        out += f"\n*...and {len(snap['files']) - 10} more files.*"

    return out

TOAST_JS = """
function showScrollToast() {
    const existing = document.getElementById('gitback-toast');
    if (existing) existing.remove();

    const toast = document.createElement('div');
    toast.id = 'gitback-toast';
    toast.innerText = '👇 Scroll down to see the snapshot';
    toast.style.cssText = `
        position: fixed;
        bottom: 32px;
        left: 50%;
        transform: translateX(-50%);
        background: #1f2937;
        color: #f9fafb;
        padding: 12px 24px;
        border-radius: 999px;
        font-size: 15px;
        font-weight: 500;
        box-shadow: 0 4px 20px rgba(0,0,0,0.35);
        z-index: 9999;
        opacity: 0;
        transition: opacity 0.3s ease;
        pointer-events: none;
        white-space: nowrap;
    `;
    document.body.appendChild(toast);
    requestAnimationFrame(() => {
        requestAnimationFrame(() => { toast.style.opacity = '1'; });
    });
    setTimeout(() => {
        toast.style.opacity = '0';
        setTimeout(() => toast.remove(), 400);
    }, 3000);
    return [];
}
"""

with gr.Blocks(title="GitBack Machine", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 🕰️ GitBack Machine
### Every GitHub commit, cached forever. Never duplicated.
Search any public repo → browse commits → click any snapshot to view it instantly.

Each snapshot is stored **exactly once**, assigned to a bucket by SHA hash. Concurrent requests are deduplicated automatically.
    """)

    with gr.Row():
        repo_input = gr.Textbox(
            placeholder="owner/repo or full GitHub URL",
            label="GitHub Repository", scale=4
        )
        search_btn = gr.Button("🔍 Search", variant="primary", scale=1)

    status = gr.Markdown("")
    state = gr.Textbox(visible=False)
    commit_list = gr.Radio(choices=[], label="📋 Commits — click to view snapshot", interactive=True)
    snapshot_view = gr.Markdown("*Search a repo and select a commit to view its snapshot.*")

    search_btn.click(search_repo, inputs=repo_input, outputs=[commit_list, status, state])
    repo_input.submit(search_repo, inputs=repo_input, outputs=[commit_list, status, state])

    commit_list.change(fn=lambda: None, inputs=[], outputs=[], js=TOAST_JS)
    commit_list.change(view_snapshot, inputs=[commit_list, state], outputs=snapshot_view)

demo.launch()