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

app = FastAPI(title="Telegram UID Search")

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

DATASET = "CodeXDevloper/MergedTgDataset"

@app.get("/")
def root():
    return {"status": "running"}

@app.get("/search")
def search_uid(uid: str = Query(...)):
    # Lazy loading — dataset sirf search ke time load hoga
    ds = load_dataset(DATASET, split="train", streaming=True)
    
    for row in ds:
        if uid == str(row.get("user_id", "")):
            return {"found": True, "data": row}
    return {"found": False}