#!/usr/bin/env python3
"""
Protection Suite API Endpoints
Provides safety checking for URLs, wallets, and tokens.
"""

import asyncio
import json
import logging
import os
import time
from collections import defaultdict
from datetime import UTC, datetime
from threading import Lock
from typing import Any
from urllib.parse import urlparse

import httpx
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel

logger = logging.getLogger("protection")


# ── Rate Limiter ─────────────────────────────────────────────────────────────
class SimpleRateLimiter:
    """Simple in-memory rate limiter for protection endpoints."""

    def __init__(self):
        self._requests: dict[str, list[float]] = defaultdict(list)
        self._lock = Lock()

    def is_allowed(self, client_id: str, limit: int = 60, window: int = 60) -> tuple:
        now = time.time()
        cutoff = now - window
        with self._lock:
            self._requests[client_id] = [t for t in self._requests[client_id] if t > cutoff]
            if len(self._requests[client_id]) >= limit:
                oldest = min(self._requests[client_id])
                retry_after = int(oldest + window - now) + 1
                return False, retry_after
            self._requests[client_id].append(now)
            return True, 0


rate_limiter = SimpleRateLimiter()


def get_client_id(request: Request) -> str:
    """Extract client identifier for rate limiting."""
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        return forwarded.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


# ── Shared HTTP Client (dependency injection pattern) ───────────────────────
_shared_http_client: httpx.AsyncClient | None = None


def get_http_client() -> httpx.AsyncClient:
    """Get shared HTTP client for connection pooling."""
    global _shared_http_client
    if _shared_http_client is None:
        _shared_http_client = httpx.AsyncClient(timeout=15.0)
    return _shared_http_client


# ── Router Definition ─────────────────────────────────────────────────────────
router = APIRouter(prefix="/api/v1/protect", tags=["protection"])


# Pydantic models for request/response
class URLCheckRequest(BaseModel):
    url: str


class WalletCheckRequest(BaseModel):
    address: str
    chain: str = "ethereum"


class TokenCheckRequest(BaseModel):
    address: str
    chain: str


class ProtectionResponse(BaseModel):
    safe: bool
    risk: str = "low"  # low, medium, high, critical
    reason: str = ""
    source: str = ""  # rag, scanner, blocklist, etc.
    risk_score: int | None = None  # 0-100 for wallet/token
    flags: list[str] = []  # for wallet: ["drainer", "scammer", ...]
    error: str | None = None


async def check_rag_for_url(url: str) -> dict[str, Any] | None:
    """
    Check RAG for known scam URL.
    """
    try:
        # Import RAG service lazily to avoid circular deps
        from app.rag_service import three_pillar_search

        # Search in known_scams and scam_patterns collections for the URL
        # Add timeout to prevent slow RAG from blocking endpoint
        results = await asyncio.wait_for(
            three_pillar_search(
                query=url,
                collections=["known_scams", "scam_patterns"],
                limit=1,
                min_similarity=0.7,  # Only high confidence matches
                use_mmr=False,  # We just want the top result
                use_kg=False,  # Disable KG for speed
                use_parent_child=False,  # Disable parent-child for speed
                use_reranker=False,  # Disable reranker for speed
            ),
            timeout=0.2,  # 0.2 second timeout for RAG check
        )
        # Check if we got any results
        if results and results.get("results"):
            # If there's at least one result with similarity >= 0.7, consider it a match
            for res in results["results"]:
                if res.get("similarity", 0) >= 0.7:
                    return {
                        "safe": False,
                        "risk": "critical",
                        "reason": f"Known scam site per RAG (similarity: {res['similarity']:.2f})",
                        "source": "rag",
                    }
    except TimeoutError:
        logger.debug(f"RAG check timed out for {url}")
    except Exception as e:
        logger.debug(f"RAG check failed for {url}: {e}")
    return None


async def check_blocklist_for_url(url: str) -> dict[str, Any] | None:
    """
    Check URL against blocklist.json domains.
    """
    try:
        blocklist_path = "/app/blocklist.json"
        if not os.path.exists(blocklist_path):
            return None
        with open(blocklist_path) as f:
            blocklist = json.load(f)
        parsed = urlparse(url)
        domain = parsed.netloc.lower()
        # Remove www prefix
        if domain.startswith("www."):
            domain = domain[4:]
        if domain in blocklist:
            return {
                "safe": False,
                "risk": "critical",
                "reason": "Domain on blocklist",
                "source": "blocklist",
            }
    except Exception as e:
        logger.debug(f"Blocklist check failed for {url}: {e}")
    return None


async def scan_url_content(url: str) -> dict[str, Any] | None:
    """
    Scan URL content for malicious patterns (simple check).
    """
    try:
        client = get_http_client()
        # Add timeout to prevent slow URL fetching from blocking endpoint
        resp = await asyncio.wait_for(
            client.get(
                str(url),
                follow_redirects=True,
                headers={"User-Agent": "RMI-Protection-Bot"},
                timeout=1.5,
            ),
            timeout=2.0,  # 2 second timeout for URL fetch
        )
        if resp.status_code != 200:
            return None
        text = resp.text.lower()
        # Simple phishing indicators
        phishing_indicators = [
            "verify your account",
            "suspend",
            "urgent action required",
            "private key",
            "seed phrase",
            "recovery phrase",
            "connect wallet",
            "claim your reward",
            "airdrop",
            "free tokens",
        ]
        for indicator in phishing_indicators:
            if indicator in text:
                return {
                    "safe": False,
                    "risk": "high",
                    "reason": f"Phishing indicator: {indicator}",
                    "source": "scanner",
                }
    except TimeoutError:
        logger.debug(f"URL content scan timed out for {url}")
    except Exception as e:
        logger.debug(f"URL content scan failed for {url}: {e}")
    return None


async def feed_url_to_rag(url: str, is_safe: bool):
    """Feed URL scan result back to RAG for learning (best effort)."""
    try:
        from app.rag_service import ingest_document

        doc = {
            "content": f"URL: {url}\\nSafety: {'safe' if is_safe else 'unsafe'}\\nScanned by RMI Protection Suite",
            "metadata": {"source": "protection_suite", "url": url, "safe": is_safe},
        }
        await ingest_document(doc)
    except Exception as e:
        logger.debug(f"Failed to feed URL to RAG: {e}")


# Endpoint implementations
@router.post("/check-url", response_model=ProtectionResponse)
async def check_url(request: URLCheckRequest, http_request: Request):
    """
    Check a URL for phishing/scam content.
    Returns: {safe: bool, risk: low|medium|high|critical, reason: str, source: rag|scanner|blocklist}
    """
    try:
        client_id = get_client_id(http_request)
        allowed, retry_after = rate_limiter.is_allowed(client_id, limit=60, window=60)

        if not allowed:
            raise HTTPException(status_code=429, detail=f"Rate limit exceeded. Retry after {retry_after} seconds.")

        url = request.url.strip()
        if not url:
            raise HTTPException(status_code=400, detail="URL is required")

        # Run checks in parallel: blocklist (fastest), RAG, content scan
        blocklist_task = asyncio.create_task(check_blocklist_for_url(url))
        rag_task = asyncio.create_task(check_rag_for_url(url))
        scan_task = asyncio.create_task(scan_url_content(url))

        # Wait for all to complete (or timeout individually inside each function)
        blocklist_result, rag_result, scan_result = await asyncio.gather(
            blocklist_task, rag_task, scan_task, return_exceptions=False
        )

        # Check results in order of preference: blocklist, rag, scan
        if blocklist_result and not isinstance(blocklist_result, Exception):
            return ProtectionResponse(**blocklist_result)

        if rag_result and not isinstance(rag_result, Exception):
            return ProtectionResponse(**rag_result)

        if scan_result and not isinstance(scan_result, Exception):
            # Feed result to RAG for future
            await feed_url_to_rag(url, is_safe=scan_result.get("safe", True))
            return ProtectionResponse(**scan_result)

        # If we get here, no threats detected
        await feed_url_to_rag(url, is_safe=True)
        return ProtectionResponse(safe=True, risk="low", reason="No threats detected", source="protection_suite")
    except Exception as e:
        logger.warning(f"URL check failed for {request.url}: {e}")
        return ProtectionResponse(
            safe=True,  # Fail safe: assume safe on error
            risk="low",
            reason="URL scan unavailable",
            source="error",
            error="service_unavailable",
        )


@router.post("/check-wallet", response_model=ProtectionResponse)
async def check_wallet(request: WalletCheckRequest, http_request: Request):
    """
    Check a wallet address for risk.
    Returns: {safe: bool, risk_score: 0-100, flags: [list], reason: str, source: ...}
    """
    try:
        client_id = get_client_id(http_request)
        allowed, retry_after = rate_limiter.is_allowed(client_id, limit=30, window=60)

        if not allowed:
            raise HTTPException(status_code=429, detail=f"Rate limit exceeded. Retry after {retry_after} seconds.")

        address = request.address.strip()
        chain = request.chain.lower().strip()
        if not address:
            raise HTTPException(status_code=400, detail="Address is required")

        # Normalize chain
        if chain not in [
            "ethereum",
            "solana",
            "base",
            "bsc",
            "polygon",
            "arbitrum",
            "optimism",
            "avalanche",
            "fantom",
        ]:
            chain = "ethereum"  # default fallback

        # Use unified scanner for wallet risk with timeout
        from app.unified_scanner import scan_wallet as wallet_scan

        # We'll use the free tier scan with a timeout of 1.8s to prevent slow scans
        try:
            result = await asyncio.wait_for(wallet_scan(address, chain), timeout=1.8)
        except TimeoutError:
            logger.warning(f"Wallet scan timed out for {address}")
            return ProtectionResponse(
                safe=True,  # Fail safe: assume safe on timeout
                risk="low",
                reason="Wallet scan timed out",
                source="error",
                error="service_unavailable",
            )

        # Extract risk score and flags
        risk_score = getattr(result, "total_risk_score", 50)  # default 50 if not found
        # Determine risk level - higher score = higher risk
        if risk_score >= 80:
            risk_level = "critical"
        elif risk_score >= 60:
            risk_level = "high"
        elif risk_score >= 40:
            risk_level = "medium"
        else:
            risk_level = "low"
        # Extract flags from risk factors (simplified)
        flags = []
        if hasattr(result, "is_labeled_scammer") and result.is_labeled_scammer:
            flags.append("scammer")
        if hasattr(result, "is_labeled_sanctioned") and result.is_labeled_sanctioned:
            flags.append("sanctioned")
        if hasattr(result, "approval_risk_score") and result.approval_risk_score > 70:
            flags.append("approval_risk")
        # Convert to dict for response - safe when risk_score < 60
        return ProtectionResponse(
            safe=risk_score < 60,  # Risk score < 60 means low risk = safe
            risk=risk_level,
            reason=f"Wallet risk score: {risk_score}",
            source="unified_scanner",
            risk_score=risk_score,
            flags=flags,
        )
    except Exception as e:
        logger.warning(f"Wallet scan failed for {request.address}: {e}")
        return ProtectionResponse(
            safe=True,  # Fail safe: assume safe on error
            risk="low",
            reason="Wallet scan unavailable",
            source="error",
            error="service_unavailable",
        )


@router.post("/check-token", response_model=ProtectionResponse)
async def check_token(request: TokenCheckRequest, http_request: Request):
    """
    Check a token address for safety.
    Returns: full safety report from scanner with protection-focused summary.
    """
    try:
        client_id = get_client_id(http_request)
        allowed, retry_after = rate_limiter.is_allowed(client_id, limit=30, window=60)

        if not allowed:
            raise HTTPException(status_code=429, detail=f"Rate limit exceeded. Retry after {retry_after} seconds.")

        address = request.address.strip()
        chain = request.chain.lower().strip()
        if not address:
            raise HTTPException(status_code=400, detail="Address is required")
        if not chain:
            raise HTTPException(status_code=400, detail="Chain is required")

        # Use token scanner with timeout to prevent slow scans
        from app.token_scanner import scan_token

        try:
            scan_result = await asyncio.wait_for(
                scan_token(address, chain, tier="free", protection_mode=True), timeout=1.8
            )
        except TimeoutError:
            logger.warning(f"Token scan timed out for {address}")
            return ProtectionResponse(
                safe=True,  # Fail safe: assume safe on timeout
                risk="low",
                reason="Token scan timed out",
                source="error",
                error="service_unavailable",
            )
        # Convert ScanResult to our response - higher safety score = safer token
        safety_score = scan_result.safety_score
        # Safety score: higher = safer, so invert for risk level
        if safety_score >= 80:
            risk_level = "low"
        elif safety_score >= 60:
            risk_level = "medium"
        elif safety_score >= 40:
            risk_level = "high"
        else:
            risk_level = "critical"
        # Extract risk flags
        flags = scan_result.risk_flags if hasattr(scan_result, "risk_flags") else []
        # Determine source
        source = "token_scanner"
        if hasattr(scan_result, "modules_run") and scan_result.modules_run:
            # Use first module that ran as source hint
            source = (
                scan_result.modules_run[0].get("module", "token_scanner")
                if scan_result.modules_run
                else "token_scanner"
            )
        # Safe when safety_score >= 60
        return ProtectionResponse(
            safe=safety_score >= 60,
            risk=risk_level,
            reason=f"Token safety score: {safety_score}",
            source=source,
            risk_score=safety_score,
            flags=flags,
        )
    except Exception as e:
        logger.warning(f"Token scan failed for {request.address}: {e}")
        return ProtectionResponse(
            safe=True,  # Fail safe
            risk="low",
            reason="Token scan unavailable",
            source="error",
            error="service_unavailable",
        )


@router.get("/domains")
async def get_domains():
    """
    Return the blocklist JSON for extension download.
    Reads from /root/browser-extension/blocklist.json, /app/blocklist.json, and RAG known scams.
    Returns: {"domains": [...], "count": N}
    """
    logger.info("=== get_domains endpoint called ===")
    try:
        # Load blocklist from browser extension file
        extension_blocklist_path = "/root/browser-extension/blocklist.json"
        backend_blocklist_path = "/app/blocklist.json"
        domains_set: set[str] = set()

        logger.info(f"Loading blocklist from {extension_blocklist_path} and {backend_blocklist_path}")

        if os.path.exists(extension_blocklist_path):
            logger.info(f"Extension blocklist file exists: {extension_blocklist_path}")
            with open(extension_blocklist_path) as f:
                extension_domains = json.load(f)
                logger.info(f"Loaded {len(extension_domains)} domains from extension blocklist")
                domains_set.update(extension_domains)
        else:
            logger.warning(f"Extension blocklist file does not exist: {extension_blocklist_path}")

        if os.path.exists(backend_blocklist_path):
            logger.info(f"Backend blocklist file exists: {backend_blocklist_path}")
            with open(backend_blocklist_path) as f:
                backend_domains = json.load(f)
                logger.info(f"Loaded {len(backend_domains)} domains from backend blocklist")
                domains_set.update(backend_domains)
        else:
            logger.warning(f"Backend blocklist file does not exist: {backend_blocklist_path}")

        # Extract domains from RAG known_scams collection
        try:
            from app.rag_service import three_pillar_search

            # Query RAG for known scams (look for URLs in content)
            rag_results = await asyncio.wait_for(
                three_pillar_search(
                    query="",  # Empty query to get all known scams (relevance based on popularity)
                    collections=["known_scams"],
                    limit=100,  # Limit to avoid slowing down endpoint
                    min_similarity=0.0,  # Get all documents regardless of similarity
                    use_mmr=False,
                    use_kg=False,
                    use_parent_child=False,
                    use_reranker=False,
                ),
                timeout=0.5,  # 0.5 second timeout for RAG query
            )
            if rag_results and rag_results.get("results"):
                for res in rag_results["results"]:
                    content = res.get("content", "")
                    # Extract URL from content (format: "URL: <url>\nSafety: ...")
                    for line in content.split("\n"):
                        if line.startswith("URL: "):
                            url = line[5:].strip()
                            if url:
                                # Parse URL to extract domain
                                try:
                                    parsed = urlparse(url)
                                    domain = parsed.netloc.lower()
                                    # Remove www prefix
                                    if domain.startswith("www."):
                                        domain = domain[4:]
                                    if domain:  # Only add non-empty domains
                                        domains_set.add(domain)
                                except Exception:
                                    pass  # Skip invalid URLs
                            break  # Only take first URL line
        except TimeoutError:
            logger.debug("RAG query for domains timed out")
        except Exception as e:
            logger.debug(f"Failed to extract domains from RAG: {e}")

        domains = sorted(domains_set)
        logger.info(f"Returning {len(domains)} unique domains (including RAG)")
        return {"domains": domains, "count": len(domains)}
    except Exception as e:
        logger.error(f"Failed to load blocklist: {e}", exc_info=True)
        # Return empty list on error to avoid breaking extensions
        return {"domains": [], "count": 0}


@router.get("/health")
async def protection_health():
    """
    Health check for protection modules.
    Returns status of RAG, blocklist, scanners, etc.
    """
    health_status = {
        "status": "ok",
        "timestamp": datetime.now(UTC).isoformat(),
        "modules": {
            "rag": "unknown",
            "blocklist": "unknown",
            "unified_scanner": "unknown",
            "token_scanner": "unknown",
        },
    }
    # Check blocklist file (backend)
    try:
        blocklist_path = "/app/blocklist.json"
        if os.path.exists(blocklist_path):
            with open(blocklist_path) as f:
                json.load(f)  # validate JSON
            health_status["modules"]["blocklist"] = "ok"
        else:
            health_status["modules"]["blocklist"] = "missing_file"
    except Exception as e:
        health_status["modules"]["blocklist"] = f"error: {e}"

    # Check RAG (simple ping)
    try:
        from app.rag_service import three_pillar_search

        # We'll do a quick query
        await three_pillar_search(query="test", collections=["known_scams"], limit=1, min_similarity=0.5)
        health_status["modules"]["rag"] = "ok"
    except Exception as e:
        health_status["modules"]["rag"] = f"error: {e}"

    # Check unified scanner (by trying to import)
    try:
        health_status["modules"]["unified_scanner"] = "ok"
    except Exception as e:
        health_status["modules"]["unified_scanner"] = f"error: {e}"

    # Check token scanner
    try:
        health_status["modules"]["token_scanner"] = "ok"
    except Exception as e:
        health_status["modules"]["token_scanner"] = f"error: {e}"

    # Overall status: if any critical module is error, set status to degraded
    critical_modules = ["rag", "blocklist", "unified_scanner", "token_scanner"]
    for mod in critical_modules:
        if health_status["modules"][mod] != "ok":
            health_status["status"] = "degraded"
            break

    return health_status
