main.py

# main.py
import os
import asyncio
from dotenv import load_dotenv
from fastapi import FastAPI, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
from datetime import datetime
from contextlib import asynccontextmanager
from supabase import create_client, Client

# Import your agent functions
from agents.security_agent import run_security_check, load_security_model, clean_email_text
from agents.summary_agent import run_summarization, load_summary_model
from agents.classification_agent import load_classification_model, run_classification
from agents.extraction_agent import run_extraction
from agents.translation_agent import run_translation 
from agents.reply_agent import generate_smart_reply

load_dotenv()
SUPABASE_URL = os.getenv("VITE_SUPABASE_URL")
SUPABASE_KEY = os.getenv("VITE_SUPABASE_ANON_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

# --- Global Lock for the Sweeper ---
is_sweeping = False

# 1. The Startup Loader
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("🚀 Booting up the PFE Intelligence Backend...")
    
    # --- DYNAMIC PATH RESOLUTION ---
    # This works automatically on both Windows (Local) and Linux (Hugging Face)
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    MODEL_PATH = os.path.join(BASE_DIR, "agents", "security_model_v2")
    # -------------------------------
    
    load_security_model(MODEL_PATH)
    load_summary_model()  
    load_classification_model()
    yield

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

# 2. Schemas
class BoundingBox(BaseModel):
    x0: int
    y0: int
    x1: int
    y1: int
    text: str

class EmailPayload(BaseModel):
    id: str
    user_id: str
    threadId: str
    sender: str
    date: datetime
    body: str
    historyId: Optional[str] = None
    subject: Optional[str] = None
    body_plain: Optional[str] = None
    snippet: Optional[str] = None

class TranslationRequest(BaseModel):
    id: str
    body: str 
    body_plain: str 
    target_lang: str = "English"

class SecurityFeedback(BaseModel):
    email_id: str
    is_secure: bool
    text: str

class ReplyRequest(BaseModel):
    email_id: str
    text: str
    sender: str
    tone: str

def update_agent_status(agent_id: str, status: str, current_task: str, increment: bool = False):
    try:
        update_data = {
            "status": status,
            "currentTask": current_task,
            "lastActivity": datetime.utcnow().isoformat()
        }
        
        if increment:
            res = supabase.table("agent_status").select("emailsProcessed").eq("id", agent_id).execute()
            if res.data:
                current_count = res.data[0].get("emailsProcessed", 0)
                update_data["emailsProcessed"] = current_count + 1

        supabase.table("agent_status").update(update_data).eq("id", agent_id).execute()
    except Exception as e:
        print(f"⚠️ Failed to update agent {agent_id}: {e}")

# 3. Core Processing Logic
async def process_single_email(email_data: EmailPayload, raw_text: str, force_safe: bool = False):
    try:
        # --- 1. GATEKEEPER (Security / mapped to agent-1) ---
        if force_safe:
            update_agent_status("agent-1", "processing", f"Human override for {email_data.sender}...")
            clean_text = clean_email_text(raw_text)
            sec_result = {
                "is_secure": True,
                "phishing_score": 0.0,
                "clean_text": clean_text
            }
            update_agent_status("agent-1", "idle", "Bypassed (Human Override)", increment=True)
        else:
            update_agent_status("agent-1", "processing", f"Scanning {email_data.sender} for threats...")
            sec_result = await run_security_check(raw_text)
            update_agent_status("agent-1", "idle", "Standing by...", increment=True)
        
        clean_text = sec_result["clean_text"]
        ai_summary, ai_category, extracted_entities, ai_urgency = None, None, None, None
        performance_metrics = {"phishing_score": sec_result["phishing_score"]}
        
        if sec_result["is_secure"]:
            # --- 3. CLASSIFIER (mapped to agent-2) ---
            update_agent_status("agent-2", "processing", "Categorizing email content...")
            class_result = await run_classification(subject=email_data.subject or "", text=clean_text)
            ai_category = class_result["category"]
            update_agent_status("agent-2", "idle", "Standing by...", increment=True)
            
            # --- 4. EXTRACTOR / PRIORITY (mapped to agent-3) ---
            update_agent_status("agent-3", "processing", "Extracting entities & urgency level...")
            ext_result = await run_extraction(clean_text, ai_category)
            extracted_entities = ext_result["extracted_data"]
            ai_urgency = ext_result["urgency"]
            update_agent_status("agent-3", "idle", "Standing by...", increment=True)

            # --- 2. PROCESSOR / SUMMARIZER (mapped to agent-4) ---
            update_agent_status("agent-4", "processing", "Generating concise summary...")
            sum_result = await run_summarization(clean_text)
            ai_summary = sum_result["summary"]
            update_agent_status("agent-4", "idle", "Standing by...", increment=True)
            
            performance_metrics.update({
                "summary": {"latency": sum_result["latency"]},
                "classification": class_result["metrics"],
                "extraction": {
                    "latency_s": ext_result["latency_s"],
                    "model": "llama-3.1-8b-instant"
                }
            })
        else:
            ai_summary = "⚠️ Processing aborted: Security threat detected."

        # 5. SAVE TO SUPABASE
        supabase.table("Email").update({
            "is_secure": sec_result["is_secure"], 
            "ai_summary": ai_summary,
            "ai_category": ai_category,
            "ai_extracted_entities": extracted_entities,
            "ai_urgency_score": ai_urgency,
            "performance_metrics": performance_metrics,
            "is_analyzed": True
        }).eq("id", email_data.id).execute()
        
        print(f"✅ Full pipeline complete for {email_data.id}")
    except Exception as e:
        print(f"⚠️ Pipeline crashed for {email_data.id}: {e}")

# 4. The Continuous Sweeper Task
async def continuous_sweeper_task():
    global is_sweeping
    if is_sweeping:
        print("⏩ Sweeper is already running in the background. Skipping new trigger.")
        return
        
    is_sweeping = True
    print("🚀 Master Sweeper Engine started...")
    
    try:
        while True:
            response = supabase.table("Email").select("*").eq("is_analyzed", False).limit(10).execute()
            pending_emails = response.data
            
            if not pending_emails:
                print("✨ All emails have been analyzed! Sweeper going to sleep.")
                break
                
            print(f"📦 Processing new batch of {len(pending_emails)} emails...")
            
            for email_row in pending_emails:
                email_payload = EmailPayload(
                    id=email_row["id"],
                    user_id=email_row["user_id"],
                    threadId=email_row["threadId"],
                    sender=email_row["sender"],
                    date=email_row["date"],
                    body=email_row["body"] or "",
                    historyId=email_row.get("historyId"),
                    subject=email_row.get("subject"),
                    body_plain=email_row.get("body_plain"),
                    snippet=email_row.get("snippet")
                )
                raw_text = email_row.get("body_plain") or email_row.get("body") or email_row.get("snippet") or ""
                await process_single_email(email_payload, raw_text)

            if len(pending_emails) < 10:
                print("✨ Final batch complete! Sweeper going to sleep.")
                break

            print("⏳ Batch complete. Resting for 15 seconds...")
            await asyncio.sleep(15) 
            
    except Exception as e:
        print(f"⚠️ Master Sweeper crashed: {e}")
    finally:
        is_sweeping = False


# 5. API Endpoints
@app.post("/api/analyze-email")
async def analyze_email(email: EmailPayload, background_tasks: BackgroundTasks):
    
    # --- NEW: Check the Feedback memory before scanning ---
    feedback_res = supabase.table("SecurityFeedback").select("is_phishing").eq("email_id", email.id).execute()
    
    force_safe = False
    if feedback_res.data and feedback_res.data[0].get("is_phishing") is False:
        force_safe = True
        print(f"🧠 Memory triggered: Bypassing security scan for {email.id} (Human Override)")

    db_insert_data = {
        "id": email.id, "user_id": email.user_id, "threadId": email.threadId,
        "historyId": email.historyId, "subject": email.subject, "sender": email.sender,
        "date": email.date.isoformat(), "body": email.body, "body_plain": email.body_plain,
        "snippet": email.snippet, "is_analyzed": False
    }
    
    supabase.table("Email").upsert(db_insert_data).execute()
    
    raw_text = email.body_plain or email.body or email.snippet or ""
    
    background_tasks.add_task(process_single_email, email, raw_text, force_safe)
    return {"status": "success", "message": "Email queued for full processing"}

@app.post("/api/process-pending")
async def process_pending_emails(background_tasks: BackgroundTasks):
    background_tasks.add_task(continuous_sweeper_task)
    return {"status": "success", "message": "Sweeper engine started in the background."}


@app.post("/api/translate-email")
async def translate_email_endpoint(payload: TranslationRequest):
    try:
        update_agent_status("agent-translation", "processing", f"Translating to {payload.target_lang}...")

        result = await run_translation(payload.body, payload.body_plain, payload.target_lang)

        if not result:
            update_agent_status("agent-translation", "idle", "Skipped - Content too short")
            return {"status": "skipped", "message": "Content too short or failed"}

        email_record = supabase.table("Email").select("performance_metrics").eq("id", payload.id).execute()
        current_metrics = email_record.data[0].get("performance_metrics", {}) if email_record.data else {}
        current_metrics["translation"] = result["metrics"]

        supabase.table("Email").update({
            "body_translated": result["translated_html"],
            "body_plain": result["refined_plain"], 
            "performance_metrics": current_metrics
        }).eq("id", payload.id).execute()

        update_agent_status("agent-translation", "idle", "Standing by...", increment=True)
        return {"status": "success", "data": result}

    except Exception as e:
        print(f"Translation Error: {e}")
        update_agent_status("agent-translation", "idle", "Error encountered")
        return {"status": "error", "message": str(e)}

@app.post("/api/feedback/security")
async def log_security_feedback(feedback: SecurityFeedback, background_tasks: BackgroundTasks):
    try:
        cleaned_text = clean_email_text(feedback.text)

        db_insert_data = {
            "email_id": feedback.email_id,
            "corrected_text": cleaned_text,
            "is_phishing": not feedback.is_secure, 
            "processed": False 
        }
        
        supabase.table("SecurityFeedback").upsert(
            db_insert_data, 
            on_conflict="email_id"
        ).execute()
        
        print(f"🛡️ Security feedback logged for email {feedback.email_id}")

        email_res = supabase.table("Email").select("*").eq("id", feedback.email_id).execute()
        
        if email_res.data:
            email_row = email_res.data[0]
            
            email_payload = EmailPayload(
                id=email_row["id"], user_id=email_row["user_id"], threadId=email_row["threadId"],
                sender=email_row["sender"], date=email_row["date"], body=email_row["body"] or "",
                historyId=email_row.get("historyId"), subject=email_row.get("subject"),
                body_plain=email_row.get("body_plain"), snippet=email_row.get("snippet")
            )
            raw_text = email_row.get("body_plain") or email_row.get("body") or email_row.get("snippet") or ""
            
            supabase.table("Email").update({
                "is_analyzed": False,
                "ai_summary": None, 
                "is_secure": True
            }).eq("id", feedback.email_id).execute()
            
            background_tasks.add_task(process_single_email, email_payload, raw_text, True)
            
        return {"status": "success", "message": "Feedback logged and email queued for full processing."}

    except Exception as e:
        print(f"⚠️ Failed to process security feedback: {e}")
        return {"status": "error", "message": str(e)}
    

@app.post("/api/generate-reply")
async def generate_reply_endpoint(payload: ReplyRequest):
    try:
        update_agent_status("agent-reply", "processing", f"Drafting a {payload.tone} reply...")
        
        draft = await generate_smart_reply(
            text=payload.text, 
            sender=payload.sender, 
            tone=payload.tone
        )
        
        update_agent_status("agent-reply", "idle", "Standing by...", increment=True)
        
        return {"status": "success", "reply": draft}
        
    except Exception as e:
        update_agent_status("agent-reply", "idle", "Error encountered")
        return {"status": "error", "message": str(e)}
@app.get("/")
def read_root():
    return {"status": "The PFE Intelligence Backend is awake and ready!"}