"""
PhishShield — Hugging Face Spaces Deployment
ML-Powered Phishing Detection for Emails and Websites
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import os
import re
import time
import sys

# Add current directory to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from rag_engine import RAGEngine
from agent import PhishingAgent
from website_scanner import WebsiteScanner

app = FastAPI(
    title="PhishShield API",
    description="ML-Powered Phishing Detection for Emails and Websites",
    version="2.0.0"
)

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

# ── Initialize Models ──────────────────────────────────────────────────────────
print("\n" + "═" * 65)
print("🛡️  PhishShield — Loading AI Models")
print("═" * 65)

print("\n📧 Loading Email Classifier (JellyPhish)...")
rag = RAGEngine()

print("\n🌐 Loading Website Scanner (URLBERT)...")
scanner = WebsiteScanner()

print("\n🧠 Loading Decision Agent...")
agent = PhishingAgent()

# Warm up models
print("\n⏳ Warming up models...")
try:
    rag.query("Test email for warming up")
    scanner.scan_url("https://example.com")
    print("✅ Models ready")
except Exception as e:
    print(f"⚠️ Warm-up warning: {e}")

print("\n" + "═" * 65)
print("✅ All engines initialized")
print("═" * 65)

# ── Static Files ───────────────────────────────────────────────────────────────
STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
if os.path.exists(STATIC_DIR):
    app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def root():
    """Serve dashboard"""
    dashboard_path = os.path.join(STATIC_DIR, "dashboard.html")
    if os.path.exists(dashboard_path):
        with open(dashboard_path, "r", encoding="utf-8") as f:
            html = f.read()
            # Fix asset paths for Hugging Face
            html = html.replace('src="dashboard.js"', 'src="/static/dashboard.js"')
            html = html.replace('href="dashboard.css"', 'href="/static/dashboard.css"')
            return html
    return "<h1>PhishShield API</h1><p>Dashboard not found</p>"


@app.get("/dashboard")
async def dashboard():
    """Serve dashboard HTML"""
    dashboard_path = os.path.join(STATIC_DIR, "dashboard.html")
    if os.path.exists(dashboard_path):
        return FileResponse(dashboard_path)
    return {"error": "Dashboard not found"}


# ── Schemas ────────────────────────────────────────────────────────────────────

class EmailPayload(BaseModel):
    subject: Optional[str] = ""
    body: str
    sender: Optional[str] = ""
    headers: Optional[Dict[str, Any]] = {}

class WebsitePayload(BaseModel):
    url: str
    threshold: Optional[float] = 0.5


# ── Heuristics ─────────────────────────────────────────────────────────────────

PHISHING_KEYWORDS = [
    "urgent", "immediately", "verify", "password", "bank",
    "account", "suspended", "click here", "login", "confirm",
    "update", "security alert", "unauthorized", "limited time",
    "act now", "winner", "prize", "free", "expire", "internship"
]

SUSPICIOUS_DOMAINS = [
    "bit.ly", "tinyurl.com", "t.co", "goo.gl",
    "phish", "fake", "secure-login", "account-verify",
    "paypal-security", "apple-id", "signin-"
]

def extract_links(text: str) -> List[str]:
    return re.findall(r"https?://\S+", text)

def is_suspicious_link(url: str) -> bool:
    url_lower = url.lower()
    return any(domain in url_lower for domain in SUSPICIOUS_DOMAINS)

def find_keywords(text: str) -> List[str]:
    text_lower = text.lower()
    return [kw for kw in PHISHING_KEYWORDS if kw in text_lower]

def score_email(keywords: List[str], suspicious_links: List[str], sender: str) -> int:
    score = 0
    score += min(len(keywords) * 1, 5)
    score += min(len(suspicious_links) * 2, 4)
    if sender and any(d in sender.lower() for d in SUSPICIOUS_DOMAINS):
        score += 1
    return min(score, 10)


# ── API Endpoints ─────────────────────────────────────────────────────────────

@app.post("/api/analyze")
async def analyze_email(payload: EmailPayload):
    """Analyze email content"""
    t0 = time.time()
    full_text = f"{payload.subject} {payload.body}"

    links = extract_links(full_text)
    suspicious_links = [l for l in links if is_suspicious_link(l)]
    keywords = find_keywords(full_text)
    heuristic_score = score_email(keywords, suspicious_links, payload.sender)

    rag_match, rag_similarity = rag.query(full_text)

    risk_level, is_phishing, reason, explanation, final_score = agent.decide(
        heuristic_score, rag_match, rag_similarity, keywords, suspicious_links
    )

    latency = round((time.time() - t0) * 1000, 2)

    return {
        "type": "email",
        "is_phishing": is_phishing,
        "risk_level": risk_level,
        "risk_score": final_score,
        "links": links,
        "suspicious_links": suspicious_links,
        "keywords_found": keywords,
        "rag_match": rag_match,
        "rag_similarity": round(rag_similarity, 3),
        "decision_reason": reason,
        "xai_explanation": explanation,
        "latency_ms": latency
    }


@app.post("/api/analyze-website")
async def analyze_website(payload: WebsitePayload):
    """Analyze website URL"""
    try:
        t0 = time.time()
        result = scanner.scan_url(payload.url, threshold=payload.threshold)
        latency = round((time.time() - t0) * 1000, 2)
        result["latency_ms"] = latency
        
        if "phishing_probability" in result:
            result["rag_similarity"] = result["phishing_probability"] / 100
        
        return {"type": "website", **result}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/api/health")
async def health():
    """Health check"""
    return {
        "status": "ok",
        "version": "2.0.0",
        "models": {
            "email": "JellyPhish (BERT)",
            "website": "URLBERT-tiny-v4"
        }
    }


@app.get("/api/models/status")
async def models_status():
    """Model status"""
    return {
        "email_model": "loaded",
        "website_model": "loaded" if scanner.use_ml else "fallback",
        "agent_ready": True
    }


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)