from fastapi import FastAPI, HTTPException, Query
import duckdb
from contextlib import asynccontextmanager

con = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global con
    con = duckdb.connect()
    con.execute("PRAGMA memory_limit='14GB'")
    con.execute("PRAGMA threads=4")
    con.execute("SET enable_object_cache=true")
    con.execute("SET preserve_insertion_order=false")
    con.execute("INSTALL httpfs; LOAD httpfs;")
    
    URLS = [
        "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main/part1.parquet",
        "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main/part2a.parquet",
        "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main/part2b_new.parquet"
    ]
    
    # Warm-up cache (load metadata)
    for url in URLS:
        con.execute(f"SELECT 1 FROM read_parquet('{url}') LIMIT 1")
    print("🟢 All parts cached! Ready.")
    yield
    con.close()

app = FastAPI(lifespan=lifespan)
VALID_KEY = "@kzr0x"

@app.get("/")
def home(key: str = Query(...)):
    if key != VALID_KEY:
        return {"status": "error", "message": "Invalid key"}
    return {"status": "success", "message": "API Live! Use /search"}

@app.get("/search")
async def search_data(
    key: str = Query(...),
    number: str = Query(None),
    aadhar: str = Query(None),
    q: str = Query(None)
):
    if key != VALID_KEY:
        return {"status": "error", "message": "Invalid key"}
    if not number and not aadhar and not q:
        raise HTTPException(status_code=400, detail="Provide number, aadhar, or q")
    
    search_term = q or number or aadhar
    
    URLS = [
        "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main/part1.parquet",
        "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main/part2a.parquet",
        "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main/part2b_new.parquet"
    ]
    
    # Build query with per-part WHERE clauses
    queries = []
    for url in URLS:
        queries.append(f"""
            SELECT * FROM read_parquet('{url}')
            WHERE phoneNumber = ? OR aadharNumber = ?
        """)
    full_query = " UNION ALL ".join(queries) + " LIMIT 10"
    
    params = [search_term, search_term] * len(URLS)
    result = con.execute(full_query, params).fetchall()
    
    if not result:
        return {
            "status": "error",
            "message": f"'{search_term}' not found in any dataset."
        }
    
    # Format response
    data_list = []
    for row in result:
        data_list.append({
            "name": row[0],
            "fathersName": row[1],
            "phoneNumber": row[2],
            "aadharNumber": row[3],
            "otherNumber": row[4],
            "address": row[5],
            "district": row[6],
            "pincode": row[7],
            "state": row[8],
            "town": row[9],
            "source": row[10]  # icmr or inddata
        })
    
    icmr_data = [d for d in data_list if d["source"] == "icmr"]
    hitek_data = [d for d in data_list if d["source"] == "inddata"]
    
    return {
        "credits": {
            "developer": "@kzr0x",
            "team": "TEAM DARK KNIGHT APIs",
            "channel": "@API_WALLAH"
        },
        "icmr_data": icmr_data,
        "hitek_data": hitek_data,
        "results": data_list  # all matches
    }