from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import pipeline
import uvicorn

app = FastAPI(title="PrivAware API", version="2.0")

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

print("Loading phishing model...")
phishing_classifier = pipeline(
    "text-classification",
    model="V3d4nt7/privaware-phishing-detector"
)

print("Loading policy model...")
policy_classifier = pipeline(
    "text-classification",
    model="V3d4nt7/privaware-policy-classifier"
)

print("Both models loaded!")

class URLRequest(BaseModel):
    url: str

class PolicyRequest(BaseModel):
    text: str

@app.get("/")
def root():
    return {"status": "ok", "message": "PrivAware API v2 running — both models live"}

@app.get("/health")
def health():
    return {"status": "healthy", "models": ["phishing-detector", "policy-classifier"]}

@app.post("/scan-url")
def scan_url(request: URLRequest):
    if not request.url:
        raise HTTPException(status_code=400, detail="URL is required")
    try:
        result = phishing_classifier(request.url[:512])[0]
        label = result["label"]
        confidence = round(result["score"] * 100, 2)
        risk_score = round(confidence) if label == "PHISHING" else round(100 - confidence)
        return {
            "url": request.url,
            "label": label,
            "confidence": confidence,
            "risk_score": risk_score,
            "is_phishing": label == "PHISHING"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/scan-policy")
def scan_policy(request: PolicyRequest):
    if not request.text:
        raise HTTPException(status_code=400, detail="Text is required")
    try:
        result = policy_classifier(request.text[:512])[0]
        label = result["label"]
        confidence = round(result["score"] * 100, 2)
        risk_map = {"SAFE": 10, "RISKY": 55, "DECEPTIVE": 90}
        risk_score = risk_map.get(label, 50)
        return {
            "label": label,
            "confidence": confidence,
            "risk_score": risk_score,
            "is_deceptive": label == "DECEPTIVE"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/scan-full")
def scan_full(request: URLRequest):
    url_result = scan_url(request)
    combined_score = url_result["risk_score"]
    return {
        "url": request.url,
        "phishing": url_result,
        "combined_risk_score": combined_score
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)