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

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

# HF 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"
WORKERS = 20        # 20 parallel workers
BATCH_SIZE = 10     # 10 batches

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")])

def search_in_file(file_path: str, uid: str):
    """Ek file mein UID search karo (worker function)"""
    try:
        ds = load_dataset(
            DATASET,
            data_files={"train": [file_path]},
            split="train",
            streaming=True
        )
        scanned = 0
        for row in ds:
            scanned += 1
            if uid == str(row.get("user_id", "")).strip():
                return {
                    "found": True,
                    "file": file_path,
                    "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"),
                    }
                }
        return {"found": False, "file": file_path, "scanned": scanned}
    except Exception as e:
        return {"found": False, "file": file_path, "error": str(e), "scanned": 0}

@app.get("/")
def root():
    return {
        "status": "running",
        "usage": "/search?uid=385167100",
        "note": f"Parallel search with {WORKERS} workers, batch size {BATCH_SIZE}, no timeout"
    }

@app.get("/search")
def search_uid(uid: str = Query(..., description="Telegram UID")):
    """
    Poore dataset mein UID dhoondho - 20 parallel workers, no timeout.
    """
    try:
        files = get_all_files()
        if not files:
            return {"error": "Koi parquet file nahi mili", "uid": uid}

        start_time = time.time()
        total_scanned = 0

        # 20 workers, batch me files process karo
        with ThreadPoolExecutor(max_workers=WORKERS) as executor:
            # batch_size ke hisaab se chunks banao
            for i in range(0, len(files), BATCH_SIZE):
                batch = files[i:i + BATCH_SIZE]
                futures = {executor.submit(search_in_file, f, uid): f for f in batch}

                for future in as_completed(futures):
                    result = future.result()
                    total_scanned += result.get("scanned", 0)

                    if result.get("found"):
                        # Baaki futures cancel karne ki koshish
                        for f in futures:
                            f.cancel()
                        return {
                            "found": True,
                            "uid": uid,
                            "scanned": total_scanned,
                            "time_taken": round(time.time() - start_time, 2),
                            "file": result.get("file"),
                            "data": result.get("data"),
                        }

        return {
            "found": False,
            "uid": uid,
            "scanned": total_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}