src/apps/app.py

import bcrypt
# Robust monkeypatch for bcrypt-passlib compatibility 
if not hasattr(bcrypt, "__about__"):
    bcrypt.__about__ = type('About', (object,), {'__version__': bcrypt.__version__})

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Response, Depends, HTTPException, status
from typing import Optional
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import uvicorn
import asyncio
import os
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from pydantic import BaseModel, EmailStr

from utils.main import RAG
from database import engine, Base, get_db
from models import User
from auth import get_password_hash, verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM, SECRET_KEY
from email_utils import send_reset_email
from datetime import timedelta
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
import uuid

from database import engine, Base, get_db, AsyncSessionLocal

# Initialize Database and Create Default Admin
async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    
    # Ensure Admin exists
    async with AsyncSessionLocal() as db:
        result = await db.execute(select(User).where(User.username == "admin"))
        admin_user = result.scalars().first()
        
        if not admin_user:
            print("Creating default admin user...")
            new_admin = User(
                username="admin",
                email="admin@lawbot.ai",
                hashed_password=get_password_hash("ADMin"),
                role="Admin"
            )
            db.add(new_admin)
            await db.commit()
            print("Admin user created successfully.")

app = FastAPI(on_startup=[init_db])

chat_history = []

# Get the base directory
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
STATIC_DIR = os.path.join(BASE_DIR, "static")
TEMPLATES_DIR = os.path.join(BASE_DIR, "templates")

# Mount static directories
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/images", StaticFiles(directory=os.path.join(STATIC_DIR, "images")), name="images")

# Initialize Jinja2 templates
templates = Jinja2Templates(directory=TEMPLATES_DIR)


# Pydantic Models for Auth
class UserCreate(BaseModel):
    username: str
    email: EmailStr
    password: str
    role: str
    mobile_number: str # Exactly 10 digits required

class UserLogin(BaseModel):
    username: str
    password: str
    role: Optional[str] = None

class ForgotPassword(BaseModel):
    email: EmailStr

class ResetPassword(BaseModel):
    token: str
    new_password: str

class Interaction(BaseModel):
    caseId: str
    query: str
    response: str
    role: str  # Role context for this interaction

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/login")

async def get_current_user(request: Request, db: AsyncSession = Depends(get_db)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    
    # Try to get token from Header first (for API calls), then Cookie (for Page navigation)
    token = None
    auth_header = request.headers.get("Authorization")
    if auth_header and auth_header.startswith("Bearer "):
        token = auth_header.split(" ")[1]
    else:
        token = request.cookies.get("access_token")
        
    if not token:
        raise credentials_exception
        
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
    except JWTError:
        raise credentials_exception
    
    result = await db.execute(select(User).where(User.username == username))
    user = result.scalars().first()
    if user is None:
        raise credentials_exception
    return user

def role_required(required_role: str):
    async def role_checker(user: User = Depends(get_current_user)):
        if user.role != required_role and user.role != "Admin":
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=f"Access denied: Requires {required_role} role (or Admin)"
            )
        return user
    return role_checker

# Auth Routes
@app.post("/api/register")
async def register(user: UserCreate, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where((User.username == user.username) | (User.email == user.email)))
    if result.scalars().first():
        raise HTTPException(status_code=400, detail="Username or Email already registered")
    
    # Validation: Exactly 10 digits
    if not user.mobile_number.isdigit() or len(user.mobile_number) != 10:
        raise HTTPException(status_code=400, detail="Mobile number must be exactly 10 digits.")
    
    hashed_password = get_password_hash(user.password)
    new_user = User(
        username=user.username, 
        email=user.email, 
        hashed_password=hashed_password, 
        role=user.role,
        mobile_number=user.mobile_number
    )
    db.add(new_user)
    await db.commit()
    return {"message": "User created successfully"}

@app.post("/api/login")
async def login(response: Response, user: UserLogin, db: AsyncSession = Depends(get_db)):
    print(f"Login attempt for username: {user.username}")
    result = await db.execute(select(User).where(User.username == user.username))
    db_user = result.scalars().first()
    
    if not db_user:
        print(f"User not found: {user.username}")
        raise HTTPException(status_code=400, detail="Incorrect username or password")
    
    print(f"User found: {db_user.username}, checking password...")
    password_valid = verify_password(user.password, db_user.hashed_password)
    print(f"Password valid: {password_valid}")
    
    if not password_valid:
        raise HTTPException(status_code=400, detail="Incorrect username or password")
    
    # Verify role if provided
    if user.role and db_user.role != user.role and db_user.role != "Admin":
        print(f"Role mismatch: user has {db_user.role}, but tried to login as {user.role}")
        raise HTTPException(status_code=403, detail=f"Access denied: This account is registered as {db_user.role}, not {user.role}")

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(data={"sub": db_user.username, "role": db_user.role}, expires_delta=access_token_expires)
    
    # Set secure cookie - use samesite="none" for Hugging Face iframes
    response.set_cookie(
        key="access_token",
        value=access_token,
        httponly=True,
        max_age=ACCESS_TOKEN_EXPIRE_MINUTES * 60,
        samesite="none",
        secure=True  # Required for Hugging Face HTTPS and samesite="none"
    )
    
    print(f"Login successful for {db_user.username}, role: {db_user.role}")
    return {
        "access_token": access_token, 
        "token_type": "bearer", 
        "role": db_user.role,
        "question_count": db_user.question_count,
        "is_admin": db_user.role == "Admin"
    }

@app.post("/api/logout")
async def logout(response: Response):
    response.delete_cookie(
        key="access_token",
        samesite="none",
        secure=True
    )
    return {"message": "Logged out successfully"}

@app.post("/api/forgot-password")
async def forgot_password(request: ForgotPassword, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where(User.email == request.email))
    user = result.scalars().first()
    
    if user:
        token = str(uuid.uuid4())
        user.reset_token = token
        await db.commit()
        try:
            await send_reset_email(user.email, token)
        except Exception as e:
            print(f"Error sending email: {e}")
            raise HTTPException(status_code=500, detail=f"Failed to send email: {str(e)}")
        
    return {"message": "If an account exists, a reset email has been sent"}

@app.post("/api/reset-password")
async def reset_password(request: ResetPassword, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where(User.reset_token == request.token))
    user = result.scalars().first()
    
    if not user:
        raise HTTPException(status_code=400, detail="Invalid or expired reset token")
    
    user.hashed_password = get_password_hash(request.new_password)
    user.reset_token = None
    await db.commit()
    return {"message": "Password reset successful"}

@app.post("/api/save-interaction")
async def save_interaction(interaction: Interaction, db: AsyncSession = Depends(get_db), current_user: User = Depends(get_current_user)):
    from models import ChatInteraction
    
    # Create new chat interaction record linked to user with role
    new_interaction = ChatInteraction(
        case_id=interaction.caseId,
        query=interaction.query,
        response=interaction.response,
        role=interaction.role,
        user_id=current_user.id
    )
    
    db.add(new_interaction)
    await db.commit()
    await db.refresh(new_interaction)
    
    print(f"Saved interaction for user {current_user.username}: ID={new_interaction.id}, CaseID={interaction.caseId}")
    return {"status": "success", "message": "Interaction saved", "id": new_interaction.id}

@app.get("/api/interactions")
async def get_interactions(role: Optional[str] = None, db: AsyncSession = Depends(get_db), current_user: User = Depends(get_current_user)):
    from models import ChatInteraction
    from sqlalchemy import func
    
    # Build base where clause
    where_clause = ChatInteraction.user_id == current_user.id
    if role:
        where_clause = (ChatInteraction.user_id == current_user.id) & (ChatInteraction.role == role)
    
    # Subquery to find the latest interaction per case_id (filtered by role if specified)
    subquery = (
        select(
            ChatInteraction.case_id,
            func.max(ChatInteraction.created_at).label("max_created")
        )
        .where(where_clause)
        .group_by(ChatInteraction.case_id)
        .subquery()
    )
    
    # Join to get query detail for the latest message in each case
    result = await db.execute(
        select(ChatInteraction)
        .join(subquery, (ChatInteraction.case_id == subquery.c.case_id) & (ChatInteraction.created_at == subquery.c.max_created))
        .order_by(ChatInteraction.created_at.desc())
        .limit(20)
    )
    interactions = result.scalars().all()
    return [{"id": i.id, "case_id": i.case_id, "query": i.query, "created_at": i.created_at} for i in interactions]

@app.get("/api/interactions/{case_id}")
async def get_conversation_thread(case_id: str, role: Optional[str] = None, db: AsyncSession = Depends(get_db), current_user: User = Depends(get_current_user)):
    from models import ChatInteraction
    
    # Build where clause with optional role filtering
    where_clause = (ChatInteraction.user_id == current_user.id) & (ChatInteraction.case_id == case_id)
    if role:
        where_clause = where_clause & (ChatInteraction.role == role)
    
    result = await db.execute(
        select(ChatInteraction)
        .where(where_clause)
        .order_by(ChatInteraction.created_at.asc())
    )
    interactions = result.scalars().all()
    if not interactions:
        raise HTTPException(status_code=404, detail="Conversation not found")
        
    return [{"query": i.query, "response": i.response, "created_at": i.created_at} for i in interactions]

@app.delete("/api/interactions/{case_id}")
async def delete_conversation(case_id: str, db: AsyncSession = Depends(get_db), current_user: User = Depends(get_current_user)):
    from models import ChatInteraction
    from sqlalchemy import delete as sqlalchemy_delete
    
    await db.execute(
        sqlalchemy_delete(ChatInteraction)
        .where((ChatInteraction.user_id == current_user.id) & (ChatInteraction.case_id == case_id))
    )
    await db.commit()
    return {"status": "success", "message": "Conversation deleted"}

@app.get("/api/user-status")
async def get_user_status(current_user: User = Depends(get_current_user)):
    return {
        "username": current_user.username,
        "role": current_user.role,
        "question_count": current_user.question_count,
        "limit": 2 if current_user.role != "Admin" else None,
        "is_admin": current_user.role == "Admin"
    }

@app.get("/api/admin/users")
async def get_all_users(db: AsyncSession = Depends(get_db), current_user: User = Depends(role_required("Admin"))):
    result = await db.execute(select(User))
    users = result.scalars().all()
    return [
        {
            "id": u.id, 
            "username": u.username, 
            "email": u.email, 
            "role": u.role, 
            "mobile_number": u.mobile_number,
            "question_count": u.question_count,
            "is_blocked": False # Placeholder if needed
        } for u in users
    ]

# Frontend Routes for Auth
@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
    return templates.TemplateResponse("login.html", {"request": request})

@app.get("/register", response_class=HTMLResponse)
async def register_page(request: Request):
    return templates.TemplateResponse("register.html", {"request": request})

@app.get("/forgot-password", response_class=HTMLResponse)
async def forgot_password_page(request: Request):
    return templates.TemplateResponse("forgot_password.html", {"request": request})

@app.get("/reset-password", response_class=HTMLResponse)
async def reset_password_page(request: Request):
    return templates.TemplateResponse("reset_password.html", {"request": request})

@app.get("/users", response_class=HTMLResponse)
@app.get("/admin/users", response_class=HTMLResponse)
async def admin_users_page(request: Request):
    # Removing API-level Depends() here so the HTML page loads and handles auth via JS fetch()
    return templates.TemplateResponse("admin_users.html", {"request": request})

# Home and Role Selection
@app.get("/", response_class=HTMLResponse)
async def role_selection(request: Request):
    return templates.TemplateResponse("roleselection.html", {"request": request, "show_roles": False})

@app.get("/role", response_class=HTMLResponse)
async def roleselection_page(request: Request):
    return templates.TemplateResponse("roleselection.html", {"request": request, "show_roles": True})

# Chatbot Pages
@app.get("/judgechatbot.html", response_class=HTMLResponse)
async def judge_chatbot(request: Request, user: User = Depends(role_required("Judge"))):
    return templates.TemplateResponse("Judgechatbot.html", {"request": request})

@app.get("/judgedashboard.html", response_class=HTMLResponse)
async def judge_dashboard(request: Request, user: User = Depends(role_required("Judge"))):
    return templates.TemplateResponse("judgedashboard.html", {"request": request})

@app.get("/viewall.html", response_class=HTMLResponse)
async def view_all(request: Request, user: User = Depends(role_required("Judge"))):
    return templates.TemplateResponse("viewall.html", {"request": request})

@app.get("/judgecalender.html", response_class=HTMLResponse)
async def judge_calender(request: Request, user: User = Depends(role_required("Judge"))):
    return templates.TemplateResponse("judgecalender.html", {"request": request})

@app.get("/advocatedashboard.html", response_class=HTMLResponse)
async def advocate_dashboard(request: Request, user: User = Depends(role_required("Advocate/Lawyer"))):
    return templates.TemplateResponse("advocatedashboard.html", {"request": request})

@app.get("/advocateresources.html", response_class=HTMLResponse)
async def advocate_resources(request: Request, user: User = Depends(role_required("Advocate/Lawyer"))):
    return templates.TemplateResponse("advocateresources.html", {"request": request})

# ========== Other Role Pages ==========
@app.get("/woman.html", response_class=HTMLResponse)
async def woman_page(request: Request, user: User = Depends(role_required("Woman"))):
    return templates.TemplateResponse("woman.html", {"request": request})

@app.get("/citizen.html", response_class=HTMLResponse)
async def citizen_page(request: Request, user: User = Depends(role_required("Citizen"))):
    return templates.TemplateResponse("citizen.html", {"request": request})

@app.get("/minor.html", response_class=HTMLResponse)
async def minor_page(request: Request, user: User = Depends(role_required("Minor"))):
    return templates.TemplateResponse("minor.html", {"request": request})

# ========== Chatbot Pages ==========
@app.get("/studentchatbot.html", response_class=HTMLResponse)
async def student_page(request: Request, user: User = Depends(role_required("Student"))):
    return templates.TemplateResponse("studentchatbot.html", {"request": request})

@app.get("/advocatechatbot.html", response_class=HTMLResponse)
async def advocatechatbot_page(request: Request):
    return templates.TemplateResponse("advocatechatbot.html", {"request": request})

@app.get("/womanchatbot.html", response_class=HTMLResponse)
async def womanchatbot_page(request: Request):
    return templates.TemplateResponse("womanchatbot.html", {"request": request})



@app.get("/safetytips.html", response_class=HTMLResponse)
async def safetytips_page(request: Request):
    return templates.TemplateResponse("safetytips.html", {"request": request})

@app.get("/resources.html", response_class=HTMLResponse)
async def resources_page(request: Request):
    return templates.TemplateResponse("resources.html", {"request": request})

@app.get("/legalrights.html", response_class=HTMLResponse)
async def legalrights_page(request: Request):
    return templates.TemplateResponse("legalrights.html", {"request": request})

@app.get("/FIR.html", response_class=HTMLResponse)
async def fir_page(request: Request):
    return templates.TemplateResponse("FIR.html", {"request": request})

@app.get("/studentdashboard.html", response_class=HTMLResponse)
async def student_page(request: Request):
    return templates.TemplateResponse("studentdashboard.html", {"request": request})

# WebSocket for Chatbot
async def stream_text_conversational(websocket: WebSocket, query: str, role: str = "General"):
    chat_limit = 10
    temp_chat = {"user": "" ,"system":""}
    temp_chat["user"] = query
    # print(f"DEBUG: stream_text_conversational executing with role={role}")
    model_response = ""
    try:
        if role == "Citizen":
            completion = RAG(query, chat_history, role=role) 
        else:
            completion = RAG(query, chat_history, role=role)
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                await websocket.send_text(chunk.choices[0].delta.content)
                await asyncio.sleep(0.01)
                model_response += chunk.choices[0].delta.content
        
        # Signal completion to frontend
        await websocket.send_text("[DONE]")
        
        # print(model_response)
        temp_chat['system']=model_response
        chat_history.append(temp_chat)
        if len(chat_history)>chat_limit:
            chat_history.pop(0)
    except Exception as e:
        error_str = str(e)
        print(f"Chatbot error: {error_str}")
        # Show a professional message — never dump raw API errors to users
        if "404" in error_str or "No endpoints" in error_str:
            user_msg = "\n\n⚠️ **Service Temporarily Unavailable**\nThe AI provider is currently unavailable. Please send your message again — the system will automatically switch to an alternative model.\n\nIf the issue persists, please try again in a moment."
        else:
            user_msg = f"\n\n⚠️ **An error occurred.** Please try sending your message again.\n\n_If the issue persists, please contact support._"
        await websocket.send_text(user_msg)
        await websocket.send_text("[DONE]")
  
@app.websocket("/conversational_chat")
async def conversational_chat(websocket: WebSocket, role: Optional[str] = None):
    await websocket.accept()
    
    # --- STRICT ROLE ENFORCEMENT ---
    # 1. Get role from Token/Cookie (Source of Truth)
    token = websocket.cookies.get("access_token")
    token_role = "General"
    user_id = None
    question_count = 0
    
    if token:
        try:
            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
            username = payload.get("sub")
            if username:
                async for db in get_db():
                    result = await db.execute(select(User).where(User.username == username))
                    user_obj = result.scalars().first()
                    if user_obj:
                        token_role = user_obj.role
                        user_id = user_obj.id
                        question_count = user_obj.question_count
                    break
        except JWTError:
            pass

    if user_id is None:
        await websocket.send_text("\n\n⚠️ Authentication Required. Please sign in to use the chatbot.")
        await websocket.close()
        return
    
    # 2. Apply Logic based on Source of Truth
    if token_role == "Admin":
        if not role:
            role = "Admin"
    else:
        role = token_role
        
    while True:
        try:
            query = await websocket.receive_text()
            
            # --- USAGE LIMIT CHECK ---
            if token_role != "Admin" and question_count >= 2:
                limit_message = (
                    "### ✨ Free usage limit reached\n\n"
                    "You've reached the free usage limit (2 questions).\n"
                    "Further access is restricted.\n\n"
                    "Please contact the administrator for extended access:\n"
                    "LinkedIn: https://www.linkedin.com/in/vishwanath77"
                )
                await websocket.send_text(limit_message)
                continue

            print(f"Query ({role}): {query}")
            await stream_text_conversational(websocket, query, role=role)
            
            # --- INCREMENT USAGE ---
            if token_role != "Admin":
                async for db in get_db():
                    result = await db.execute(select(User).where(User.id == user_id))
                    user_to_update = result.scalars().first()
                    if user_to_update:
                        user_to_update.question_count += 1
                        question_count = user_to_update.question_count # Keep local sync
                        await db.commit()
                    break
                    
        except WebSocketDisconnect:
            chat_history.clear()
            break
        except Exception as e:
            print(f"WebSocket error: {e}")
            try:
                await websocket.send_text(
                    "\n\n⚠️ **Connection interrupted.** Please try sending your message again."
                )
                await websocket.send_text("[DONE]")
            except:
                break


# Run the application
if __name__ == "__main__":
    print("Starting Law Bot Server...")
    port = int(os.getenv("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)