"""
api.py
======
FastAPI — Phishing URL Detection API
Supports:
  - POST /predict/all    → run all 3 models (serialized via asyncio.Lock)
  - POST /predict/single → run a selected model
API key auth is read from .env (API_KEYS=key1,key2,...).
All requests and results are logged to ./logs/.
"""

from __future__ import annotations

import asyncio
import json
import logging
import math
import os
import re
import time
import uuid
#import resource
from collections import Counter
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import joblib
import pandas as pd
import psutil
from dotenv import load_dotenv
from fastapi import Depends, FastAPI, HTTPException, Request, Security, status
from fastapi.security.api_key import APIKeyHeader
from pydantic import BaseModel, Field

# =============================================================================
# BOOTSTRAP
# =============================================================================

load_dotenv()

LOG_DIR = Path("logs")
LOG_DIR.mkdir(exist_ok=True)

# ── File logger (one rotating file per day) ───────────────────────────────
log_filename = LOG_DIR / f"api_{datetime.now(timezone.utc).strftime('%Y-%m-%d')}.log"
file_handler = logging.FileHandler(log_filename, encoding="utf-8")
file_handler.setFormatter(
    logging.Formatter("%(asctime)s | %(levelname)s | %(message)s", datefmt="%Y-%m-%dT%H:%M:%S")
)

console_handler = logging.StreamHandler()
console_handler.setFormatter(logging.Formatter("%(asctime)s | %(levelname)s | %(message)s"))

logger = logging.getLogger("phishing_api")
logger.setLevel(logging.INFO)
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# =============================================================================
# API KEY AUTH
# =============================================================================

_raw_keys = os.getenv("API_KEYS", "")
VALID_API_KEYS: set[str] = {k.strip() for k in _raw_keys.split(",") if k.strip()}

if not VALID_API_KEYS:
    logger.warning("No API_KEYS found in .env — all requests will be rejected!")

API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)


def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> str:
    if not api_key or api_key not in VALID_API_KEYS:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing API key.",
        )
    return api_key


# =============================================================================
# CONSTANTS  (kept in sync with the three inference scripts)
# =============================================================================

PHISHING_KEYWORDS = [
    "login", "verify", "update", "secure", "account",
    "bank", "confirm", "signin", "password", "reset", "admin",
]
SUSPICIOUS_TLDS = ["xyz", "top", "club", "tk", "ml", "ga", "cf", "buzz", "cn", "ru"]
SHORTENERS      = ["bit.ly", "tinyurl.com", "t.co", "goo.gl", "is.gd", "ow.ly", "buff.ly"]
TRUSTED_DOMAINS = ["github.com", "youtube.com", "microsoft.com", "apple.com"]
BRAND_REAL_DOMAINS: dict[str, list[str]] = {
    "paypal":    ["paypal.com"],
    "google":    ["google.com", "google.co"],
    "facebook":  ["facebook.com", "fb.com"],
    "bank":      [],
    "apple":     ["apple.com"],
    "microsoft": ["microsoft.com", "live.com", "outlook.com"],
    "amazon":    ["amazon.com", "aws.amazon.com"],
}

# =============================================================================
# FEATURE HELPERS
# =============================================================================

def _is_ip(domain: str) -> bool:
    return bool(re.match(r"^\d{1,3}(\.\d{1,3}){3}$", domain))


def _entropy(s: str) -> float:
    if not s:
        return 0.0
    prob = [n / len(s) for _, n in Counter(s).items()]
    return -sum(p * math.log2(p) for p in prob)


def _brand_impersonation(url_lower: str, domain: str) -> int:
    for brand, real_domains in BRAND_REAL_DOMAINS.items():
        if brand in url_lower:
            if any(domain == rd or domain.endswith("." + rd) for rd in real_domains):
                continue
            return 1
    return 0


def _kw_domain(domain: str) -> int:
    return sum(k in domain for k in PHISHING_KEYWORDS)


def _kw_path(path: str) -> int:
    return sum(k in path for k in PHISHING_KEYWORDS)


def _is_trusted(domain: str, path: str) -> bool:
    for trusted in TRUSTED_DOMAINS:
        if domain == trusted or domain.endswith("." + trusted):
            if _kw_path(path) == 0:
                return True
    return False


def extract_features(url: str) -> dict[str, Any]:
    url = url.strip()
    if not url.startswith(("http://", "https://")):
        url = "https://" + url

    parsed      = urlparse(url)
    domain      = parsed.netloc.lower().split(":")[0]
    path        = parsed.path.lower()
    query       = parsed.query.lower()
    length      = len(url)
    num_digits  = sum(c.isdigit() for c in url)
    num_special = sum(not c.isalnum() for c in url)
    tld         = domain.split(".")[-1] if "." in domain else ""

    return {
        "url_length":                   length,
        "domain_length":                len(domain),
        "path_length":                  len(path),
        "num_subdomains":               domain.count("."),
        "num_hyphens_url":              url.count("-"),
        "num_hyphens_domain":           domain.count("-"),
        "num_digits":                   num_digits,
        "digit_ratio":                  num_digits / length if length > 0 else 0,
        "num_special_chars":            num_special,
        "special_char_ratio":           num_special / length if length > 0 else 0,
        "num_query_params":             query.count("&") + (1 if query else 0),
        "num_phishing_keywords_domain": _kw_domain(domain),
        "num_phishing_keywords_path":   _kw_path(path),
        "is_http_only":                 int(url.startswith("http://") and not url.startswith("https://")),
        "https_in_domain":              int("https" in domain),
        "has_www":                      int("www." in domain),
        "tld_length":                   len(tld),
        "is_ip_address":                int(_is_ip(domain)),
        "is_suspicious_tld":            int(tld in SUSPICIOUS_TLDS),
        "is_shortener":                 int(any(s in domain for s in SHORTENERS)),
        "has_brand_name":               _brand_impersonation(url.lower(), domain),
        "has_at_symbol":                int("@" in url),
        "has_double_slash_path":        int("//" in path),
        "domain_entropy":               _entropy(domain),
    }


# =============================================================================
# MODEL REGISTRY
# =============================================================================

class AlgorithmName(str, Enum):
    lightgbm           = "lightgbm"
    decision_tree      = "decision_tree"
    logistic_regression = "logistic_regression"


MODEL_CONFIG: dict[AlgorithmName, dict[str, str]] = {
    AlgorithmName.lightgbm:            {"model": "phishing_model.pkl",    "features": "features.pkl"},
    AlgorithmName.decision_tree:       {"model": "phishing_model_dt.pkl", "features": "features_dt.pkl"},
    AlgorithmName.logistic_regression: {"model": "phishing_model_lr.pkl", "features": "features_lr.pkl"},
}

# Lazy-loaded model cache  {algorithm: (model_object, feature_list)}
_model_cache: dict[AlgorithmName, tuple[Any, list[str]]] = {}


def _load_model(algorithm: AlgorithmName) -> tuple[Any, list[str]]:
    if algorithm not in _model_cache:
        cfg = MODEL_CONFIG[algorithm]
        if not Path(cfg["model"]).exists():
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail=f"Model file '{cfg['model']}' not found for algorithm '{algorithm.value}'.",
            )
        if not Path(cfg["features"]).exists():
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail=f"Features file '{cfg['features']}' not found for algorithm '{algorithm.value}'.",
            )
        model    = joblib.load(cfg["model"])
        features = joblib.load(cfg["features"])
        _model_cache[algorithm] = (model, features)
        logger.info("Loaded model: %s", algorithm.value)
    return _model_cache[algorithm]

def get_memory_usage_mb() -> float:
    """Returns the current memory usage of the process in MB."""
    # resource.getrusage(resource.RUSAGE_SELF).ru_maxrss returns bytes on Linux
    # but kilobytes on macOS. On cPanel (Linux), it will be in kilobytes.
    usage_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return usage_kb / 1024.0


# =============================================================================
# CORE PREDICTION (synchronous — called inside thread pool)
# =============================================================================

def _predict(url: str, algorithm: AlgorithmName, threshold: float = 0.5) -> dict[str, Any]:
    _process = psutil.Process(os.getpid())

    url_to_parse = url if url.startswith(("http://", "https://")) else "https://" + url
    parsed = urlparse(url_to_parse)
    domain = parsed.netloc.lower().split(":")[0]
    path   = parsed.path.lower()

    # Whitelist fast-path
    if _is_trusted(domain, path):
        mem = _process.memory_info().rss / (1024 * 1024)
        return {
            "algorithm":             algorithm.value,
            "label":                 "LEGITIMATE",
            "whitelisted":           True,
            "phishing_probability":  0.0,
            "threshold":             threshold,
            "timings_ms": {
                "feature_extraction": 0.0,
                "model_inference":    0.0,
                "total":              0.0,
            },
            "memory_mb": {
                "before":      round(mem, 4),
                "after":       round(mem, 4),
                "fluctuation": 0.0,
            },
        }

    total_t0 = time.perf_counter()
    mem_before = _process.memory_info().rss / (1024 * 1024)
    #mem_before = get_memory_usage_mb()

    # Feature extraction
    feat_t0  = time.perf_counter()
    features = extract_features(url)
    feat_t1  = time.perf_counter()

    model, feature_list = _load_model(algorithm)

    # Fill any missing features with 0
    for f in feature_list:
        features.setdefault(f, 0)

    X = pd.DataFrame([features])[feature_list]

    # Inference
    pred_t0 = time.perf_counter()
    prob    = float(model.predict_proba(X)[0][1])
    pred_t1 = time.perf_counter()

    total_t1   = time.perf_counter()
    mem_after  = _process.memory_info().rss / (1024 * 1024)
    #mem_after = get_memory_usage_mb()

    label = "PHISHING" if prob >= threshold else "LEGITIMATE"

    return {
        "algorithm":             algorithm.value,
        "label":                 label,
        "whitelisted":           False,
        "phishing_probability":  round(prob, 6),
        "threshold":             threshold,
        "timings_ms": {
            "feature_extraction": round((feat_t1 - feat_t0) * 1000, 4),
            "model_inference":    round((pred_t1 - pred_t0) * 1000, 4),
            "total":              round((total_t1 - total_t0) * 1000, 4),
        },
        "memory_mb": {
            "before":      round(mem_before, 4),
            "after":       round(mem_after, 4),
            "fluctuation": round(mem_after - mem_before, 4),
        },
    }


# =============================================================================
# SERIALIZATION LOCK  (one inference at a time across all workers)
# =============================================================================

_inference_lock = asyncio.Lock()


async def _safe_predict(url: str, algorithm: AlgorithmName, threshold: float) -> dict[str, Any]:
    """Wraps _predict under the global lock so requests are serialized."""
    async with _inference_lock:
        # Run blocking joblib/sklearn/lgbm call in thread pool
        loop   = asyncio.get_running_loop()
        result = await loop.run_in_executor(None, _predict, url, algorithm, threshold)
    return result


# =============================================================================
# PYDANTIC SCHEMAS
# =============================================================================

class SinglePredictRequest(BaseModel):
    url:       str            = Field(..., example="https://paypal-secure-update.com/login")
    algorithm: AlgorithmName  = Field(..., example="lightgbm")
    threshold: float          = Field(default=0.5, ge=0.0, le=1.0)


class AllPredictRequest(BaseModel):
    url:       str   = Field(..., example="https://paypal-secure-update.com/login")
    threshold: float = Field(default=0.5, ge=0.0, le=1.0)


# =============================================================================
# LOG HELPER
# =============================================================================

def _log_request_result(
    request_id: str,
    endpoint: str,
    client_ip: str | None,
    payload: dict,
    result: dict,
    elapsed_ms: float,
) -> None:
    entry = {
        "request_id": request_id,
        "timestamp":  datetime.now(timezone.utc).isoformat(),
        "endpoint":   endpoint,
        "client_ip":  client_ip,
        "payload":    payload,
        "result":     result,
        "total_request_ms": round(elapsed_ms, 4),
    }
    logger.info("REQUEST | %s", json.dumps(entry))


# =============================================================================
# APP
# =============================================================================

app = FastAPI(
    title="Phishing URL Detection API",
    description=(
        "Detect phishing URLs using Decision Tree, LightGBM, or Logistic Regression.\n\n"
        "Authenticate with the `X-API-Key` header."
    ),
    version="1.0.0",
)


@app.get("/health", tags=["Utility"])
async def health():
    """Simple liveness check — no auth required."""
    return {"status": "ok", "timestamp": datetime.now(timezone.utc).isoformat()}


# ---------------------------------------------------------------------------
# POST /predict/all
# ---------------------------------------------------------------------------

@app.post("/predict/all", tags=["Prediction"])
async def predict_all(
    body:    AllPredictRequest,
    request: Request,
    _key:    str = Depends(verify_api_key),
):
    """
    Run the URL through **all three** models.
    Requests are serialized — each model runs sequentially under a lock.
    """
    request_id = str(uuid.uuid4())
    wall_t0    = time.perf_counter()

    results: list[dict] = []
    errors:  list[dict] = []

    for algo in AlgorithmName:
        try:
            res = await _safe_predict(body.url, algo, body.threshold)
            results.append(res)
        except HTTPException as exc:
            errors.append({"algorithm": algo.value, "error": exc.detail})
        except Exception as exc:
            errors.append({"algorithm": algo.value, "error": str(exc)})

    wall_t1    = time.perf_counter()
    elapsed_ms = (wall_t1 - wall_t0) * 1000

    # Aggregate verdict: PHISHING if ANY model says PHISHING
    labels    = [r["label"] for r in results]
    aggregate = "PHISHING" if "PHISHING" in labels else "LEGITIMATE"

    response = {
        "request_id":        request_id,
        "url":               body.url,
        "threshold":         body.threshold,
        "aggregate_verdict": aggregate,
        "results":           results,
        "errors":            errors,
        "total_request_ms":  round(elapsed_ms, 4),
    }

    _log_request_result(
        request_id=request_id,
        endpoint="/predict/all",
        client_ip=request.client.host if request.client else None,
        payload=body.model_dump(),
        result=response,
        elapsed_ms=elapsed_ms,
    )

    return response


# ---------------------------------------------------------------------------
# POST /predict/single
# ---------------------------------------------------------------------------

@app.post("/predict/single", tags=["Prediction"])
async def predict_single(
    body:    SinglePredictRequest,
    request: Request,
    _key:    str = Depends(verify_api_key),
):
    """
    Run the URL through a **single selected** model.
    """
    request_id = str(uuid.uuid4())
    wall_t0    = time.perf_counter()

    try:
        result = await _safe_predict(body.url, body.algorithm, body.threshold)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))

    wall_t1    = time.perf_counter()
    elapsed_ms = (wall_t1 - wall_t0) * 1000

    response = {
        "request_id":       request_id,
        "url":              body.url,
        "threshold":        body.threshold,
        **result,
        "total_request_ms": round(elapsed_ms, 4),
    }

    _log_request_result(
        request_id=request_id,
        endpoint="/predict/single",
        client_ip=request.client.host if request.client else None,
        payload=body.model_dump(),
        result=response,
        elapsed_ms=elapsed_ms,
    )

    return response


# =============================================================================
# ENTRYPOINT
# =============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=False)