main.py

# VERSION: v5.13.15 (SymPy Coordinate Fix + UI Prep)
# RESTART TRIGGER: 2026-03-16T23:05:00
from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, Header
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from sse_starlette.sse import EventSourceResponse
from fastapi.templating import Jinja2Templates
from typing import Optional, List
import logging
import base64
import json
import io
import sys
import os
import asyncio
from datetime import datetime
from pydantic import BaseModel
from typing import Any

class AskQuestionRequest(BaseModel):
    context_data: dict | Any
    question: str
    student_name: str = "תלמיד"
    id_token: Optional[str] = None

# --- HEALTH CHECK : Top-level Dependency Verification ---
# We do this before standard imports to ensure a clear error message 
# if the virtual environment is inactive and libraries are missing.
try:
    import cv2
    import numpy as np
except ModuleNotFoundError as e:
    print(f"🔥 [HEALTH-CHECK FAILED] Missing critical dependency: {e}. Are you running inside the .venv?")
    sys.exit(1)

from orchestrator import orchestrator, build_standard_response
from quota_system import quota_manager
from quota_system_v2 import quota_manager_v2
from config import IS_PRODUCTION, ENV
from firebase_manager import firebase_manager
from utils.image_processor import enhance_image_for_math_ocr

if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8')

# הגדרת לוגר HamoraServer
try:
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler("/tmp/server.log", encoding="utf-8"),
            logging.StreamHandler(sys.stdout)
        ]
    )
except PermissionError:
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[logging.StreamHandler(sys.stdout)]
    )

logger = logging.getLogger("HamoraServer")

# --- INFRA HARDENING: Global Async Exception Handler ---
def custom_async_exception_handler(loop, context):
    """
    Catches unhandled asynchronous exceptions to prevent the Event Loop from crashing.
    """
    msg = context.get("exception", context["message"])
    logger.critical(f"🚨 [ASYNC-CRASH-PREVENTION] Caught unhandled exception in Event Loop: {msg}")
    # The loop remains alive. We just log the critical error.

# --- INFRA HARDENING: Health Check Function ---
def verify_system_health():
    """
    Verifies execution environment and critical dependencies.
    """
    logger.info("🩺 [HEALTH-CHECK] Verifying core dependencies and environment...")
    
    # Check if running in a virtual environment
    if sys.prefix == sys.base_prefix:
        logger.warning("⚠️ [HEALTH-CHECK] Not running inside a virtual environment (.venv). Proceeding anyway...")
        
    logger.info(f"✅ [HEALTH-CHECK] cv2 version: {cv2.__version__}, numpy: {np.__version__}")
    logger.info(f"✅ [HEALTH-CHECK] Environment: {ENV.upper()}, Production Mode: {IS_PRODUCTION}")
    
    # V5.8.1: API Key Validation
    if not os.environ.get("GOOGLE_API_KEY"):
        logger.error("❌ [HEALTH-CHECK] GOOGLE_API_KEY is missing! Gemini calls will fail.")
    else:
        logger.info("✅ [HEALTH-CHECK] GOOGLE_API_KEY is detected.")



# --- INFRA HARDENING: Lifespan Context Manager ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup Phase
    verify_system_health()
    
    # V5.13.16: HARD STARTUP - Ensure Firebase is initialized before routes accept traffic
    logger.info("🚀 [STARTUP] Initializing Firebase SDK (Hard Start)...")
    firebase_manager.initialize()
    
    # Register Global Async Exception Handler
    loop = asyncio.get_running_loop()
    loop.set_exception_handler(custom_async_exception_handler)
    logger.info("🛡️ [STARTUP] Global Async Exception Handler registered.")
    

    
    yield # Yield control back to FastAPI
    
    # Shutdown Phase
    logger.info("🛑 [SHUTDOWN] BuddyMath Server is shutting down cleanly.")


# Application Setup
app = FastAPI(title="BuddyMath Server - OpenCV Engine", lifespan=lifespan)

# Setup Jinja templates
base_dir = os.path.dirname(os.path.abspath(__file__))
templates = Jinja2Templates(directory=os.path.join(base_dir, "templates"))

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Static files for audio fallback
os.makedirs("/tmp/static", exist_ok=True)
app.mount("/static", StaticFiles(directory="/tmp/static"), name="static")

@app.get("/")
async def root():
    return {"status": f"BuddyMath API V5.10.1 ({ENV.upper()})", "engine": "OpenCV Base + Security Hardening"}

async def verify_admin_access(request: Request) -> Optional[str]:
    """
    Centralized admin verification logic.
    Returns UID if access is granted, otherwise raises HTTPException.
    """
    auth_header = request.headers.get('Authorization')
    if not auth_header or not auth_header.startswith('Bearer '):
        raise HTTPException(status_code=401, detail="Unauthorized: Missing Token")

    token = auth_header.split('Bearer ')[1].strip()
    uid = None
    
    from config import DEV_BYPASS_TOKEN
    if not IS_PRODUCTION and token == DEV_BYPASS_TOKEN:
        uid = "dev-bypass-user"
        logger.info("🛠️ [ADMIN-AUTH] Using DEV Auth Bypass Token.")
    else:
        decoded = firebase_manager.verify_token(token)
        if not decoded:
            raise HTTPException(status_code=401, detail="Unauthorized: Invalid Token")
        uid = decoded.get('uid')

    if not uid:
        raise HTTPException(status_code=401, detail="Unauthorized: Invalid UID")

    # Strict Firestore Role Check
    try:
        db = firebase_manager.get_db()
        user_doc = db.collection('users').document(uid).get()
        if not user_doc.exists and uid != "dev-bypass-user":
             raise HTTPException(status_code=403, detail="Forbidden: User document missing")
        
        user_data = user_doc.to_dict() if user_doc.exists else {'role': 'admin', 'isAdmin': True}
        if user_data.get('role') != 'admin' and not user_data.get('isAdmin'):
             logger.warning(f"🚨 [ADMIN-AUTH] Unauthorized attempt by UID: {uid}")
             raise HTTPException(status_code=403, detail="Forbidden: Admin access required")
        
        return uid
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"❌ [ADMIN-AUTH] Database error: {e}")
        raise HTTPException(status_code=500, detail="Internal server error during auth")

@app.get("/admin/stats")
async def get_admin_stats(request: Request):
    """
    V5.10.1: Returns usage and cost statistics.
    Enforced strict Admin authorization.
    """
    await verify_admin_access(request)

    from cost_tracker import LOG_FILE, PRICING
    stats = {
        "total_input_tokens": 0,
        "total_output_tokens": 0,
        "total_cost_usd": 0.0,
        "request_count": 0,
        "total_users": 0
    }
    
    # 1. Count users using optimized query
    try:
        db = firebase_manager.get_db()
        results = db.collection('users').count().get()
        stats["total_users"] = results[0][0].value
    except Exception as e:
        logger.error(f"Error counting users: {e}")

    # 2. Parse usage logs
    if os.path.exists(LOG_FILE):
        try:
            with open(LOG_FILE, "r", encoding="utf-8") as f:
                for line in f:
                    try:
                        entry = json.loads(line)
                        stats["total_input_tokens"] += entry.get("input_tokens", 0)
                        stats["total_output_tokens"] += entry.get("output_tokens", 0)
                        stats["request_count"] += 1
                    except: continue
            
            # Calculate total cost
            stats["total_cost_usd"] = (stats["total_input_tokens"] / 1e6 * PRICING["input"]) + \
                                    (stats["total_output_tokens"] / 1e6 * PRICING["output"])
        except Exception as e:
            logger.error(f"Error parsing log file: {e}")
            
    return {"status": "success", "cost_summary": stats}

class QuotaUpdateRequest(BaseModel):
    uid: str
    daily_limit: Optional[int] = None
    monthly_budget: Optional[int] = None
    total_purchased: Optional[int] = None
    is_unlimited: Optional[bool] = None # V6.0

@app.post("/admin/update_quota")
async def update_quota(request: Request, req: QuotaUpdateRequest):
    """
    V5.10.1: Updates user quota.
    """
    await verify_admin_access(request)
    try:
        db = firebase_manager.get_db()
        update_data = {}
        if req.daily_limit is not None:
             update_data['quota_limit'] = req.daily_limit
        if req.monthly_budget is not None:
             update_data['monthly_token_budget'] = req.monthly_budget
        if req.total_purchased is not None:
             update_data['wallet.total_purchased_tokens'] = req.total_purchased
        if req.is_unlimited is not None:
             update_data['is_unlimited'] = req.is_unlimited
        
        if not update_data:
             return {"status": "error", "message": "No data to update"}
             
        db.collection('users').document(req.uid).update(update_data)
        logger.info(f"📊 [ADMIN] Updated quota for {req.uid}: {update_data}")
        return {"status": "success"}
    except Exception as e:
        logger.error(f"Failed to update quota: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/admin/fix_balance/{uid}")
async def fix_balance(uid: str, request: Request):
    """
    V5.14.7: Temporary fix for negative token balances.
    """
    await verify_admin_access(request)
    try:
        db = firebase_manager.get_db()
        user_ref = db.collection('users').document(uid)
        doc = user_ref.get()
        if doc.exists:
            data = doc.to_dict()
            wallet = data.get('wallet', {})
            balance = wallet.get('token_balance', 0)
            if balance < 0:
                pos_balance = abs(balance)
                user_ref.update({'wallet.token_balance': pos_balance})
                logger.info(f"✅ [ADMIN] Fixed negative balance for {uid}: {balance} -> {pos_balance}")
                return {"status": "success", "fixed_to": pos_balance}
        return {"status": "no_fix_needed"}
    except Exception as e:
        logger.error(f"Failed to fix balance: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/admin/reset_usage/{uid}")
async def reset_usage(uid: str, request: Request):
    """
    V5.10.1: Resets monthly usage to 0.
    """
    await verify_admin_access(request)
    try:
        db = firebase_manager.get_db()
        db.collection('users').document(uid).update({'used_tokens_this_month': 0})
        logger.info(f"🔄 [ADMIN] Reset usage for {uid}")
        return {"status": "success"}
    except Exception as e:
        logger.error(f"Failed to reset usage: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/admin/clear_devices/{uid}")
async def clear_devices(uid: str, request: Request):
    """
    V5.10.1: Clears all allowed devices for a user.
    """
    await verify_admin_access(request)
    try:
        db = firebase_manager.get_db()
        db.collection('users').document(uid).update({
            'allowed_devices': [], 
            'masterDeviceId': firestore.DELETE_FIELD
        })
        logger.info(f"📱 [ADMIN] Cleared devices for {uid}")
        return {"status": "success"}
    except Exception as e:
        logger.error(f"Failed to clear devices: {e}")
        raise HTTPException(status_code=500, detail=str(e))

class StatusUpdateRequest(BaseModel):
    status: str

@app.post("/admin/update_user_status/{uid}")
async def update_user_status(uid: str, req: StatusUpdateRequest, request: Request):
    """
    V5.10.1: Updates user status (approved/blocked/etc).
    """
    await verify_admin_access(request)
    try:
        db = firebase_manager.get_db()
        db.collection('users').document(uid).update({'status': req.status})
        logger.info(f"🛡️ [ADMIN] Updated status for {uid} to {req.status}")
        return {"status": "success"}
    except Exception as e:
        logger.error(f"Failed to update status: {e}")
        raise HTTPException(status_code=500, detail=str(e))



@app.post("/admin/delete_user/{uid}")
async def delete_user_full(uid: str, request: Request):
    """
    V5.10.1: Deletes user from both Firestore and Firebase Auth.
    Enforced strict Admin authorization.
    """
    await verify_admin_access(request)

    try:
        # 1. Delete from Firestore
        db = firebase_manager.get_db()
        db.collection('users').document(uid).delete()
        
        # 2. Delete from Firebase Auth
        try:
            from firebase_admin import auth
            auth.delete_user(uid)
            logger.info(f"🗑️ [ADMIN] User {uid} deleted from Auth and Firestore.")
        except auth.UserNotFoundError:
            logger.warning(f"⚠️ [ADMIN] User {uid} not found in Auth, but deleted from Firestore.")
        except Exception as e:
            logger.error(f"❌ [ADMIN] Auth deletion failed for {uid}: {e}")
            # We continue because the DB part is done

        return {"status": "success", "message": f"User {uid} removed."}
    except Exception as e:
        logger.error(f"❌ [ADMIN] Deletion failed: {e}")
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.post("/solve_stream")
async def solve_stream(
    user: Optional[str] = Form(None),
    student_name: Optional[str] = Form(None),
    grade: str = Form("י'"),
    student_gender: str = Form("M"),
    mode: str = Form("solve"),
    user_note: Optional[str] = Form(None),
    file: UploadFile = File(...)
):
    """
    V5.8.0: המורה למתמטיקה - Multipart & OpenCV Base.
    מקבל קובץ ישירות מהפלאטר ומפענח אותו עם OpenCV.
    """
    final_student_name = student_name or user or "תלמיד"
    uid = None
    print(f"🚀 🟢 BIT-LOG: Received Multipart request from {final_student_name}. Grade: {grade}")
    
    # Quota Check
    if final_student_name == "dev-bypass-user":
        is_allowed, msg, current_usage, limit = True, "Dev Bypass", 0, 999
    else:
        is_allowed, msg, current_usage, limit = quota_manager.check_limit(final_student_name)
    
    if not is_allowed:
        response_content = build_standard_response(
            final_answer=f"הגעת למכסה היומית ({limit} שאלות)",
            teacher_summary="נא להמתין למחר לקבלת מכסה חדשה.",
            logic_error=True,
            response_type="error"
        )
        response_content["error"] = "QUOTA_EXCEEDED" 
        return JSONResponse(status_code=429, content=response_content)
    quota_manager.increment_usage(final_student_name)

    try:
        # 1. קריאת הבינארי
        original_bytes = await file.read()
        print(f"📸 [BIT-LOG] Image received. Size: {len(original_bytes)} bytes")
        
        # 1b. V5.13.14: Pre-processing Layer
        image_bytes = enhance_image_for_math_ocr(original_bytes)
        print(f"📸 [BIT-LOG] Image enhanced. Size: {len(image_bytes)} bytes")

        # 2. OpenCV Decoder
        nparr = np.frombuffer(image_bytes, np.uint8)
        img_cv2 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img_cv2 is None:
            print("❌ [BIT-LOG] OpenCV failed to decode image!")
            raise HTTPException(status_code=400, detail="Invalid image data")

        print(f"✅ [BIT-LOG] OpenCV Matrix Ready: {img_cv2.shape}")

        # 3. OCR & Solving Pipeline (Streaming)
        print("🚀 [TRACE-MAIN] Initiating streaming orchestrator.solve_problem...")

        async def event_generator():
            try:
                async for event in orchestrator.solve_problem(
                    problem_text="", # Will be extracted by OCR
                    grade=grade,
                    student_name=final_student_name,
                    student_gender=student_gender,
                    user_note=user_note,
                    image_data=image_bytes,
                    mode=mode,
                    uid=uid
                ):
                    # SSE Protocol: yield a dict with "data" key
                    yield {
                        "event": "message",
                        "id": event.question_id,
                        "data": event.model_dump_json() # Pydantic v2
                    }
            except Exception as e:
                logger.error(f"STREAMING ERROR: {e}")
                yield {
                    "event": "error",
                    "data": json.dumps({"error": str(e)})
                }

        return EventSourceResponse(event_generator())

    except Exception as e:
        logger.exception("CRITICAL FLOW ERROR")
        print(f"🔥 [BIT-LOG] CRITICAL ERROR: {str(e)}")
        import traceback
        traceback.print_exc()
        response_content = build_standard_response(
            final_answer="שגיאה בפענוח התמונה או התרגיל",
            teacher_summary="המורה למתמטיקה מתנצל, אך חלה שגיאה לא צפויה.",
            logic_error=True,
            response_type="error"
        )
        return JSONResponse(status_code=500, content=response_content)

@app.post("/v2/solve_stream")
async def solve_stream_v2(
    request: Request,
    user: Optional[str] = Form(None),
    student_name: Optional[str] = Form(None),
    grade: str = Form("י'"),
    student_gender: str = Form("M"),
    mode: str = Form("solve"),
    user_note: Optional[str] = Form(None),
    session_id: Optional[str] = Form(None), # V318.0: Tutor Session Support
    files: List[UploadFile] = File(...)
):
    """
    V2: Token-based Auth and Firestore Quota Management.
    """
    auth_header = request.headers.get('Authorization')
    if not auth_header or not auth_header.startswith('Bearer '):
        logger.warning("🚨 [V2_ENDPOINT] Missing or invalid Authorization header.")
        return JSONResponse(status_code=401, content={"error": "Unauthorized: Missing Token"})

    id_token = auth_header.split('Bearer ')[1].strip()
    
    uid = None
    from config import IS_PRODUCTION, DEV_BYPASS_TOKEN
    
    # V5.9.3: Log incoming token for debugging
    logger.info(f"🔑 [V2_ENDPOINT] Received Token (len={len(id_token)}): {id_token[:5]}...")

    # V5.9.1: DEV Auth Bypass (Strict non-prod check)
    if not IS_PRODUCTION and id_token == DEV_BYPASS_TOKEN:
        logger.info("🛠️ [V2_ENDPOINT] Using DEV Auth Bypass Token.")
        uid = "dev-bypass-user"
    else:
        decoded_token = firebase_manager.verify_token(id_token)
        if not decoded_token:
            logger.warning("🚨 [V2_ENDPOINT] Invalid or expired Firebase ID token.")
            return JSONResponse(status_code=401, content={"error": "Unauthorized: Invalid Token"})
        uid = decoded_token.get('uid')

    final_student_name = student_name or user or "תלמיד"
    print(f"🚀 🟢 [V2] Received request from UID: {uid} ({final_student_name}). Grade: {grade}")
    device_id = request.headers.get('Device-ID')
    
    # V5.10.0: Fetch user tier for Digital Binder (History) support
    user_tier = "student_basic"
    try:
        db = firebase_manager.get_db()
        user_doc = db.collection('users').document(uid).get()
        if user_doc.exists:
            user_tier = user_doc.to_dict().get('tier', 'student_basic')
    except Exception as e:
        logger.error(f"Error fetching user tier: {e}")

    # V2 Quota Check (Firestore)
    if uid == "dev-bypass-user":
        is_allowed, msg, current_usage, limit = True, "Dev Bypass", 0, 999
    else:
        is_allowed, msg, current_usage, limit = quota_manager_v2.check_limit(uid, device_id=device_id)
    
    if not is_allowed:
        response_content = build_standard_response(
            final_answer=f"הגעת למכסה היומית ({limit} שאלות)",
            teacher_summary="נא להמתין למחר לקבלת מכסה חדשה או לשדרג לפרימיום.",
            logic_error=True,
            response_type="error"
        )
        response_content["error"] = "QUOTA_EXCEEDED" 
        return JSONResponse(status_code=403, content=response_content) # Changed to 403 Forbidden for quota specifically

    # V5.15.0: Pencil Economy Pre-flight Check (Wait for at least 2,000 tokens)
    has_pencils, balance = quota_manager_v2.check_wallet(uid, min_required=2000)
    if not has_pencils:
        response_content = build_standard_response(
            final_answer="לא נותרו לך מספיק עפרונות לביצוע הפעולה.",
            teacher_summary="העיפרון הושחז עד הסוף! נא לרכוש חבילת עפרונות חדשה.",
            logic_error=True,
            response_type="error"
        )
        response_content["error"] = "PAYMENT_REQUIRED"
        response_content["token_balance"] = balance
        return JSONResponse(status_code=402, content=response_content)

    # Only increment usage if OCR/Solving process starts successfully
    
    try:
        # V316.5: Sort incoming files by filename to ensure image_00, image_01... order
        files.sort(key=lambda x: x.filename)
        
        # 1. קריאת הבינארי
        original_bytes_list = []
        for single_file in files:
            original_bytes_list.append(await single_file.read())
            
        print(f"📸 [V2-LOG] Received {len(original_bytes_list)} images.")
        
        # 1b. V5.13.14: Pre-processing Layer (Map over all images)
        image_bytes_list = [enhance_image_for_math_ocr(b) for b in original_bytes_list]
        print(f"📸 [V2-LOG] Images enhanced for OCR Accuracy.")

        if not image_bytes_list:
            raise HTTPException(status_code=400, detail="No images provided")

        # 2. OpenCV Decoder (validate the first image)
        nparr = np.frombuffer(image_bytes_list[0], np.uint8)
        img_cv2 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img_cv2 is None:
            print("❌ [V2-LOG] OpenCV failed to decode first image!")
            raise HTTPException(status_code=400, detail="Invalid image data")

        print(f"✅ [V2-LOG] OpenCV Matrix Ready: {img_cv2.shape}")

        # Increment quota AFTER we are sure the image is valid
        quota_manager_v2.increment_usage(uid)

        # 3. OCR & Solving Pipeline (Streaming)
        print("🚀 [TRACE-V2] Initiating streaming orchestrator.solve_problem with multiple images...")

        async def event_generator():
            import cost_tracker
            cost_tracker.current_request_tokens.set(0)
            try:
                # orchestrator handles the rest using student_name for pedagogical stuff, but quota is already handled.
                async for event in orchestrator.solve_problem(
                    problem_text="", # Will be extracted by OCR
                    grade=grade,
                    student_name=final_student_name,
                    student_gender=student_gender,
                    user_note=user_note,
                    image_data_list=image_bytes_list,
                    mode=mode,
                    uid=uid,
                    session_id=session_id, # V318.0: Pass through
                    tier=user_tier # V5.10.0: Pass tier for history saving
                ):
                    # SSE Protocol: yield a dict with "data" key
                    yield {
                        "event": "message",
                        "id": event.question_id,
                        "data": event.model_dump_json() # Pydantic v2
                    }
            except Exception as e:
                logger.error(f"STREAMING ERROR (V2): {e}")
                yield {
                    "event": "error",
                    "data": json.dumps({"error": str(e)})
                }
            finally:
                # V5.10.0: Ensure token deduction even on crash/disconnect
                total_tokens = 0
                try:
                    total_tokens = cost_tracker.current_request_tokens.get()
                except Exception:
                    pass
                
                if total_tokens > 0:
                    try:
                        quota_manager_v2.increment_usage(uid, increment_questions=0, tokens_used=total_tokens)
                        print(f"🪙 [V2-QUOTA] Deducted {total_tokens} tokens for UID: {uid}")
                    except ValueError as ve:
                        # V5.15.0: Log overdraft attempt that was blocked by atomic transaction
                        logger.error(f"❌ [V2-QUOTA] Atomic Overdraft Blocked: {ve}")
                        # This avoids the connection hanging if the finally block crashes

        return EventSourceResponse(event_generator())

    except Exception as e:
        logger.exception("CRITICAL FLOW ERROR (V2)")
        print(f"🔥 [V2-LOG] CRITICAL ERROR: {str(e)}")
        import traceback
        traceback.print_exc()
        response_content = build_standard_response(
            final_answer="שגיאה בפענוח התמונה או התרגיל",
            teacher_summary="המורה למתמטיקה מתנצל, אך חלה שגיאה לא צפויה.",
            logic_error=True,
            response_type="error"
        )
        return JSONResponse(status_code=500, content=response_content)

@app.post("/explain_step")
async def explain_step(request: Request):
    data = await request.json()
    
    # Auth Hardening (V3.1)
    auth_header = request.headers.get('Authorization')
    token = (data.get("id_token") or (auth_header.split('Bearer ')[1] if auth_header and auth_header.startswith('Bearer ') else None))
    
    if not token:
        return JSONResponse(status_code=401, content={"error": "Unauthorized: Missing Token"})
        
    decoded_token = firebase_manager.verify_token(token)
    if not decoded_token:
        return JSONResponse(status_code=401, content={"error": "Unauthorized: Invalid Token"})
    
    uid = decoded_token.get('uid')
    
    # Quota check (Gate only, no increment for simple explanations yet)
    is_allowed, msg, _, _ = quota_manager_v2.check_limit(uid)
    if not is_allowed:
        return JSONResponse(status_code=403, content={"error": "QUOTA_EXCEEDED", "message": msg})
        
    res = await orchestrator.explain_specific_step(data.get("context"), data.get("step_text"), data.get("student_name"))
    return JSONResponse(content=res)

@app.post("/ask_question")
async def ask_question(request: Request, ask_req: AskQuestionRequest):
    data = ask_req.dict()
    
    # Auth Hardening (V3.1)
    auth_header = request.headers.get('Authorization')
    token = (data.get("id_token") or (auth_header.split('Bearer ')[1] if auth_header and auth_header.startswith('Bearer ') else None))
    
    if not token:
        return JSONResponse(status_code=401, content={"error": "Unauthorized: Missing Token"})
        
    decoded_token = firebase_manager.verify_token(token)
    if not decoded_token:
        return JSONResponse(status_code=401, content={"error": "Unauthorized: Invalid Token"})
    
    uid = decoded_token.get('uid')
    
    # Quota check
    is_allowed, msg, _, _ = quota_manager_v2.check_limit(uid)
    if not is_allowed:
        return JSONResponse(status_code=403, content={"error": "QUOTA_EXCEEDED", "message": msg})

    res = await orchestrator.ask_question(data.get("context_data"), data.get("question"), data.get("student_name"))
    return JSONResponse(content=res)
    
@app.get("/pay", response_class=HTMLResponse)
async def payment_page(request: Request, uid: Optional[str] = None):
    """
    Serves the Premium Payment Web Page.
    Fetches the student name from Firestore if uid is provided.
    """
    student_name = ""
    if uid:
        try:
            db = firebase_manager.get_db()
            user_doc = db.collection('users').document(uid).get()
            if user_doc.exists:
                student_name = user_doc.to_dict().get("student_name", "")
        except Exception as e:
            logger.error(f"Failed to fetch user for payment page: {e}")
            
    return templates.TemplateResponse("payment.html", {
        "request": request, 
        "uid": uid or "", 
        "student_name": student_name
    })

class UpgradeRequest(BaseModel):
    uid: str
    parent_email: str

@app.post("/api/upgrade_success")
async def upgrade_success(req: UpgradeRequest):
    """
    Mock Webhook for successful payment. 
    Updates the Firestore user to Premium Tier.
    """
    if not req.uid:
        raise HTTPException(status_code=400, detail="Missing user ID")
        
    try:
        db = firebase_manager.get_db()
        user_ref = db.collection('users').document(req.uid)
        user_ref.set({
            "tier": "parent_premium",
            "parent_email": req.parent_email,
            "monthly_token_budget": 2800000
        }, merge=True)
        
        # V5.10.0: Add 20 Pencils (340k tokens) with Debt Absorption
        quota_manager_v2.add_tokens_with_absorption(req.uid, 340000)
        
        logger.info(f"🎉 UPGRADED USER {req.uid} to parent_premium")
        return {"status": "success", "message": "User upgraded successfully"}
    except Exception as e:
        logger.error(f"Failed to upgrade user {req.uid}: {e}")
        raise HTTPException(status_code=500, detail="Database update failed")
    
class MigrationRequest(BaseModel):
    batch_size: Optional[int] = 50

@app.post("/admin/migrate_to_v2")
async def migrate_to_v2(request: Request, req: MigrationRequest):
    """
    V280.2: Admin endpoint to migrate users to V2 Quota system.
    """
    await verify_admin_access(request)
        
    try:
        from scripts.migrate_users_to_cloud import migrate_users
        import asyncio
        
        # Run migration in background so we don't block
        asyncio.create_task(asyncio.to_thread(migrate_users))
        
        return {"status": "success", "message": "Migration started in background."}
    except Exception as e:
        logger.error(f"Migration error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

class ReportRequest(BaseModel):
    uid: str
    student_name: str
    parent_email: str
    week_id: str

@app.post("/v2/send_weekly_report")
async def send_weekly_report(req: ReportRequest):
    """
    Generates and emails the weekly AI Assessment report to the parent.
    """
    try:
        import os
        from report_generator import report_generator
        
        # 1. Produce HTML
        html_content = report_generator.produce_weekly_report(
            uid=req.uid,
            week_id=req.week_id,
            student_name=req.student_name
        )
        
        # 2. Generate PDF
        pdf_path = f"/tmp/report_{req.uid}_{req.week_id}.pdf"
        # Create tmp dir if it doesn't exist (local dev)
        os.makedirs(os.path.dirname(pdf_path), exist_ok=True)
        report_generator.export_to_pdf(html_content, pdf_path)
        
        # 3. Email PDF
        success = report_generator.send_report_email(
            parent_email=req.parent_email,
            student_name=req.student_name,
            pdf_path=pdf_path
        )
        
        # Cleanup
        try:
            if os.path.exists(pdf_path):
                os.remove(pdf_path)
        except Exception as e:
            print(f"Cleanup failed for {pdf_path}: {e}")
            
        if success:
            return {"status": "success", "message": f"Report sent to {req.parent_email}"}
            
        return JSONResponse(status_code=500, content={"status": "error", "message": "Failed to send email via SendGrid."})
        
    except Exception as e:
        logger.error(f"Failed to generate report: {e}")
        import traceback
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)