from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from datasets import load_dataset
import os

os.environ["HF_HOME"] = "/tmp/hf"

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

DATASET = "CodeXDevloper/MergedTgDataset"

@app.get("/")
def root():
    return {"status": "running", "usage": "/search?uid=385167100"}

@app.get("/search")
def search_uid(uid: str = Query(...)):
    try:
        ds = load_dataset(DATASET, split="train", streaming=True)
        scanned = 0
        for row in ds:
            scanned += 1
            if uid == str(row.get("user_id", "")).strip():
                return {"found": True, "uid": uid, "data": row}
            if scanned >= 5000000:
                break
        return {"found": False, "uid": uid, "scanned": scanned}
    except Exception as e:
        return {"error": str(e)}