from fastapi import FastAPI
import sqlite3
import os
import shutil
from huggingface_hub import hf_hub_download

app = FastAPI()
DATASET_REPO = "sheeturt/boat-master-data"
DB_NAME = "boat_master.db"

def join_db():
    # Agar DB pehle se join ho chuka hai toh skip karo
    if os.path.exists(DB_NAME) and os.path.getsize(DB_NAME) > 2000000000:
        return
    
    print("🛠 Downloading and Joining DB parts...")
    try:
        with open(DB_NAME, 'wb') as master:
            for i in range(1, 5): # Check all 4 parts
                part_path = hf_hub_download(repo_id=DATASET_REPO, filename=f"boat_master.db.part{i}", repo_type="dataset")
                with open(part_path, 'rb') as part:
                    shutil.copyfileobj(part, master) # Efficient copying
        print("✅ Join Complete!")
    except Exception as e:
        print(f"❌ Join Error: {e}")

@app.on_event("startup")
async def startup_event():
    join_db()

@app.get("/search")
def search(num: str = None, email: str = None):
    try:
        if not os.path.exists(DB_NAME):
            return {"error": "Database files are being downloaded. Try in 5 minutes."}
            
        conn = sqlite3.connect(DB_NAME)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        if num:
            # Query with LIMIT to prevent memory overflow
            cursor.execute("SELECT phone, email, raw_info FROM boat_data WHERE phone LIKE ? LIMIT 50", (f"%{num}%",))
        elif email:
            cursor.execute("SELECT phone, email, raw_info FROM boat_data WHERE email = ? LIMIT 50", (email,))
        else:
            return {"error": "Please provide 'num' or 'email' parameter."}
            
        res = [dict(row) for row in cursor.fetchall()]
        conn.close()
        return {"count": len(res), "results": res}
    except Exception as e:
        return {"error": str(e)}

@app.get("/")
def home():
    return {"status": "Running", "message": "Use /search?num=123 or /search?email=abc"}