from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from datasets import load_dataset
from huggingface_hub import HfApi, login
import os
import time

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

# HF_TOKEN se login (fast + authenticated)
HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN:
    login(token=HF_TOKEN)

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

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

DATASET = "CodeXDevloper/MergedTgDataset"
PART_SIZE = 100000  # har part mein 1 lakh rows
MAX_PARTS = 50      # ek request mein max 50 parts
TIME_LIMIT = 25     # 25 second ke andar ruk jao

def get_all_files():
    api = HfApi()
    files = api.list_repo_files(DATASET, repo_type="dataset")
    return sorted([f for f in files if f.endswith(".parquet")])

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

@app.get("/search")
def search_uid(uid: str = Query(..., description="Telegram UID")):
    """
    Poore dataset mein UID dhoondho - parts mein divided, fast.
    """
    try:
        files = get_all_files()
        
        if not files:
            return {"error": "Koi parquet file nahi mili"}
        
        # Saari files ko parts mein load karo
        ds = load_dataset(
            DATASET,
            data_files={"train": files},
            split="train",
            streaming=True
        )
        
        start_time = time.time()
        scanned = 0
        
        for row in ds:
            scanned += 1
            
            # UID match karo
            if uid == str(row.get("user_id", "")).strip():
                return {
                    "found": True,
                    "uid": uid,
                    "scanned": scanned,
                    "time_taken": round(time.time() - start_time, 2),
                    "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"),
                    }
                }
            
            # Time limit check (har 5000 rows)
            if scanned % 5000 == 0:
                if time.time() - start_time > TIME_LIMIT:
                    return {
                        "found": False,
                        "uid": uid,
                        "scanned": scanned,
                        "time_taken": round(time.time() - start_time, 2),
                        "message": f"{TIME_LIMIT} second mein nahi mila. Phir try karo."
                    }
        
        return {
            "found": False,
            "uid": uid,
            "scanned": scanned,
            "time_taken": round(time.time() - start_time, 2),
            "message": "Poore dataset mein nahi mila"
        }
    
    except Exception as e:
        return {"error": str(e), "uid": uid}