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

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

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

# Cache directory (Docker mein writable jagah)
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["HF_DATASETS_CACHE"] = "/tmp/hf/datasets"

DATASET = "CodeXDevloper/MergedTgDataset"

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

@app.get("/search")
def search_uid(uid: str = Query(..., description="Telegram UID")):
    """Poore dataset mein UID dhoondho"""
    try:
        ds = load_dataset(DATASET, split="train", streaming=True)
        
        scanned = 0
        max_scan = 5000000  # 50 lakh rows tak
        
        for row in ds:
            scanned += 1
            
            # user_id match karo
            if uid == str(row.get("user_id", "")).strip():
                return {
                    "found": True,
                    "uid": uid,
                    "scanned": scanned,
                    "data": {
                        "user_id": row.get("user_id"),
                        "username": row.get("username"),
                        "first_name": row.get("first_name"),
                        "last_name": row.get("last_name"),
                        "phone": row.get("phone"),
                        "email": row.get("email"),
                        "status": row.get("status"),
                        "linked_id": row.get("linked_id"),
                        "linked_name": row.get("linked_name"),
                        "linked_handle": row.get("linked_handle"),
                    }
                }
            
            if scanned >= max_scan:
                break
        
        return {
            "found": False,
            "uid": uid,
            "scanned": scanned,
            "message": f"UID '{uid}' nahi mila ({scanned} rows scan kiye)"
        }
    
    except Exception as e:
        return {"error": str(e), "uid": uid}