import json
import random
import os
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import uvicorn
import io

app = FastAPI(title="HitekFresh User Data API", description="Search and stream user data from 83GB JSON", version="1.0")

JSON_FILE_PATH = "/app/users_data.json"  # Docker mein path
SAMPLE_SIZE = 100  # Random sample size

# ----- Helper: Load JSON (streaming for large files) -----
def stream_json(file_path: str):
    """Generator to stream JSON array from file."""
    with open(file_path, 'r', encoding='utf-8') as f:
        # Assume JSON is an array of objects
        data = json.load(f)
        for item in data:
            yield item

# ----- In-memory cache for random data (first 10k records) -----
_cache = None
def load_cache():
    global _cache
    if _cache is None:
        try:
            with open(JSON_FILE_PATH, 'r', encoding='utf-8') as f:
                _cache = json.load(f)[:10000]  # Load first 10k for speed
        except Exception as e:
            _cache = []
    return _cache

# ----- Models -----
class UserData(BaseModel):
    mobile: str
    name: str
    fname: Optional[str] = None
    address: Optional[str] = None
    alt: Optional[str] = None
    circle: Optional[str] = None
    id: Optional[str] = None
    email: Optional[str] = None

# ----- Endpoints -----

@app.get("/")
async def root():
    return {"message": "HitekFresh API is live!", "docs": "/docs"}

@app.get("/search")
async def search_users(
    mobile: Optional[str] = Query(None, description="Search by mobile number"),
    name: Optional[str] = Query(None, description="Search by name (partial)"),
    email: Optional[str] = Query(None, description="Search by email (partial)"),
    id: Optional[str] = Query(None, description="Search by Aadhar ID"),
    limit: int = Query(10, ge=1, le=100, description="Max results")
):
    """Search users by mobile, name, email, or Aadhar ID."""
    if not any([mobile, name, email, id]):
        raise HTTPException(status_code=400, detail="Provide at least one search parameter (mobile, name, email, id)")
    
    results = []
    for user in load_cache():  # Using cache for speed
        if mobile and user.get("mobile") != mobile:
            continue
        if name and name.lower() not in user.get("name", "").lower():
            continue
        if email and email.lower() not in user.get("email", "").lower():
            continue
        if id and user.get("id") != id:
            continue
        results.append(user)
        if len(results) >= limit:
            break
    
    if not results:
        return JSONResponse({"message": "No users found", "results": []})
    return {"count": len(results), "results": results}

@app.get("/mobile/{mobile}")
async def get_by_mobile(mobile: str):
    """Get user by exact mobile number."""
    for user in load_cache():
        if user.get("mobile") == mobile:
            return user
    raise HTTPException(status_code=404, detail="User not found")

@app.get("/aadhar/{aadhar_id}")
async def get_by_aadhar(aadhar_id: str):
    """Get user by Aadhar ID."""
    for user in load_cache():
        if user.get("id") == aadhar_id:
            return user
    raise HTTPException(status_code=404, detail="User not found")

@app.get("/random")
async def get_random_users(count: int = Query(5, ge=1, le=50)):
    """Get random users from the dataset."""
    cache = load_cache()
    if not cache:
        raise HTTPException(status_code=500, detail="Data not loaded")
    sample = random.sample(cache, min(count, len(cache)))
    return {"count": len(sample), "results": sample}

@app.get("/stream-all")
async def stream_all_users():
    """Stream all users as JSON lines (NDJSON) for large downloads."""
    def generate():
        with open(JSON_FILE_PATH, 'r', encoding='utf-8') as f:
            data = json.load(f)
            for user in data:
                yield json.dumps(user) + "\n"
    return StreamingResponse(generate(), media_type="application/x-ndjson")

@app.get("/stats")
async def get_stats():
    """Get basic stats about the dataset."""
    cache = load_cache()
    return {
        "total_records_loaded": len(cache),
        "full_file_size_gb": 83.4,
        "fields": list(cache[0].keys()) if cache else []
    }

# ----- Run -----
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)