import os
import time
import logging
import asyncio
import bisect
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import snapshot_download
import pyarrow.dataset as ds
import pyarrow.compute as pc
import pyarrow.parquet as pq

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

REPO_ID = "sauravsingh2111/Tgdata"
CACHE_DIR = "/data/tgdb_cache"

dataset = None
user_id_index = None
is_ready = False
init_error = None
stats = {"startup_time": 0, "total_files": 0, "total_row_groups": 0, "total_rows": 0, "queries": 0}

ALL_COLS = ["user_id", "username", "first_name", "last_name", "phone", "email", "status", "linked_id", "linked_name", "linked_handle"]

def find_parquet(base):
    files = []
    for root, _, names in os.walk(base):
        for n in names:
            if n.endswith(".parquet"):
                files.append(os.path.join(root, n))
    return sorted(files)

class UserIdIndex:
    def __init__(self):
        self.entries = []
        self.keys = []

    def build(self, arrow_dataset):
        logger.info("Building user_id index from parquet metadata...")
        t0 = time.time()
        schema = arrow_dataset.schema
        uid_idx = schema.get_field_index("user_id")
        if uid_idx < 0:
            logger.error("user_id column not found!")
            return
        for frag in arrow_dataset.get_fragments():
            meta = frag.metadata
            path = frag.path
            for rg_idx in range(meta.num_row_groups):
                col = meta.row_group(rg_idx).column(uid_idx)
                s = col.statistics
                if s and s.min is not None and s.max is not None:
                    self.entries.append((int(s.min), int(s.max), path, rg_idx))
        self.entries.sort(key=lambda x: x[0])
        self.keys = [e[0] for e in self.entries]
        elapsed = time.time() - t0
        logger.info(f"Index built: {len(self.entries)} row groups in {elapsed:.1f}s")

    def find(self, user_id):
        idx = bisect.bisect_right(self.keys, user_id) - 1
        if idx >= 0 and idx < len(self.entries):
            mn, mx, path, rg = self.entries[idx]
            if mn <= user_id <= mx:
                return path, rg
        return None

def init_dataset():
    global dataset, user_id_index, is_ready, init_error

    try:
        os.makedirs(CACHE_DIR, exist_ok=True)
        files = find_parquet(CACHE_DIR)

        if not files:
            logger.info("Downloading dataset (~8GB)...")
            t0 = time.time()
            snapshot_download(
                repo_id=REPO_ID,
                repo_type="dataset",
                local_dir=CACHE_DIR,
                local_dir_use_symlinks=False,
            )
            logger.info(f"Download completed in {time.time()-t0:.1f}s")
            files = find_parquet(CACHE_DIR)

        logger.info(f"Creating dataset from {len(files)} files...")
        d = ds.dataset(files, format="parquet")

        rg_count = sum(f.metadata.num_row_groups for f in d.get_fragments())
        row_count = sum(f.count_rows() for f in d.get_fragments())

        idx = UserIdIndex()
        idx.build(d)

        stats["total_files"] = len(files)
        stats["total_row_groups"] = rg_count
        stats["total_rows"] = row_count
        stats["startup_time"] = time.time()

        dataset = d
        user_id_index = idx
        is_ready = True

        logger.info(f"Ready: {len(files)} files, {rg_count} row groups, {row_count:,} rows")
    except Exception as e:
        init_error = str(e)
        logger.error(f"Init failed: {e}")

async def background_init():
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, init_dataset)

@asynccontextmanager
async def lifespan(app):
    asyncio.create_task(background_init())
    yield

app = FastAPI(title="Telegram DB API", description="859M Telegram users searchable by ID", version="1.0.0", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])

def require_ready():
    if not is_ready:
        if init_error:
            raise HTTPException(status_code=503, detail=f"Init error: {init_error}")
        raise HTTPException(status_code=503, detail="Dataset loading, please retry in a moment")

def query_user(user_id: int):
    stats["queries"] += 1
    t0 = time.time()

    result = user_id_index.find(user_id)
    if result:
        path, rg_idx = result
        try:
            pf = pq.ParquetFile(path)
            tbl = pf.read_row_group(rg_idx, columns=ALL_COLS)
            tbl = tbl.filter(pc.equal(pc.field("user_id"), user_id))
            if len(tbl):
                return tbl.to_pylist()[0], time.time() - t0
        except Exception as e:
            logger.warning(f"Index lookup failed for {user_id}: {e}")

    tbl = dataset.to_table(filter=pc.equal(pc.field("user_id"), user_id), columns=ALL_COLS)
    elapsed = time.time() - t0
    return (tbl.to_pylist()[0], elapsed) if len(tbl) else (None, elapsed)

def search_users(params: dict, limit: int = 10):
    stats["queries"] += 1
    t0 = time.time()

    conditions = []
    for key in ("username", "phone", "email", "first_name", "last_name"):
        val = params.get(key)
        if val:
            conditions.append(pc.equal(pc.field(key), val))

    if not conditions:
        return [], 0, time.time() - t0

    combined = conditions[0]
    for c in conditions[1:]:
        combined = combined & c

    tbl = dataset.to_table(filter=combined, columns=["user_id", "username", "first_name", "last_name", "phone", "email", "status"])
    elapsed = time.time() - t0
    count = len(tbl)
    if not count:
        return [], 0, elapsed
    return tbl.to_pylist()[:limit], count, elapsed

@app.get("/")
async def root():
    return {
        "name": "Telegram DB API",
        "dataset": REPO_ID,
        "ready": is_ready,
        "endpoints": {
            "/user/{user_id}": "Get full details by Telegram user ID",
            "/search": "Search by username, phone, email, first_name, last_name",
            "/health": "System health & stats"
        }
    }

@app.get("/user/{user_id}")
async def get_user(user_id: int):
    require_ready()
    user, elapsed = query_user(user_id)
    if user is None:
        raise HTTPException(status_code=404, detail=f"User {user_id} not found")
    return {"found": True, "query_time_ms": round(elapsed * 1000, 2), "user": user}

@app.get("/search")
async def search(
    username: str = Query(None),
    phone: str = Query(None),
    email: str = Query(None),
    first_name: str = Query(None),
    last_name: str = Query(None),
    limit: int = Query(10, ge=1, le=100),
):
    require_ready()
    params = {k: v for k, v in {"username": username, "phone": phone, "email": email, "first_name": first_name, "last_name": last_name}.items() if v}
    if not params:
        raise HTTPException(status_code=400, detail="Provide at least one search parameter")

    rows, total, elapsed = search_users(params, limit)
    if not rows:
        raise HTTPException(status_code=404, detail="No users found")
    return {"found": True, "count": total, "returned": len(rows), "query_time_ms": round(elapsed * 1000, 2), "users": rows}

@app.get("/health")
async def health():
    uptime = round(time.time() - stats["startup_time"], 1) if stats["startup_time"] else 0
    return {
        "status": "loading" if not is_ready else ("error" if init_error else "ok"),
        "ready": is_ready,
        "error": init_error,
        "dataset": REPO_ID,
        "cached": os.path.exists(CACHE_DIR),
        "files": stats["total_files"],
        "row_groups": stats["total_row_groups"],
        "rows": stats["total_rows"],
        "queries_served": stats["queries"],
        "uptime_s": uptime
    }