"""
SmartPPT Builder — FastAPI Backend for HuggingFace Spaces.

This backend receives Python code, runs it in an isolated environment,
and returns the generated .pptx file.

Deploy to: https://huggingface.co/spaces/SmartPPTBuilder/SmartPPTbuilder1
"""

import os
import shutil
import subprocess
import sys
import tempfile
import ast
import hashlib
import json
from pathlib import Path
from typing import Optional

import firebase_admin
from firebase_admin import credentials, auth as firebase_auth

from fastapi import FastAPI, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

# ── App Setup ────────────────────────────────────────────────────────────

app = FastAPI(
    title="SmartPPT Builder API",
    description="Generate PowerPoint presentations from Python code",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://smartpptbuilder.web.app",
        "https://smartpptbuilder.firebaseapp.com",
        "http://localhost:5173",  # Dev server
        "http://localhost:4173",  # Preview server
    ],
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

# ── Firebase Admin Setup ──────────────────────────────────────────────────
# Requires a FIREBASE_SERVICE_ACCOUNT JSON string in Hugging Face Secrets
firebase_cred_json = os.environ.get("FIREBASE_SERVICE_ACCOUNT")
if firebase_cred_json:
    try:
        cred_dict = json.loads(firebase_cred_json)
        cred = credentials.Certificate(cred_dict)
        firebase_admin.initialize_app(cred)
    except Exception as e:
        print(f"Failed to initialize Firebase Admin: {e}")
else:
    print("WARNING: FIREBASE_SERVICE_ACCOUNT not set in environment. Token verification will fail.")

# ── Config ───────────────────────────────────────────────────────────────

# Approved packages that can be imported
ALLOWED_PACKAGES = {
    "pptx": "python-pptx",
    "PIL": "Pillow",
    "cv2": "opencv-python-headless",
    "sklearn": "scikit-learn",
    "skimage": "scikit-image",
    "yaml": "pyyaml",
    "bs4": "beautifulsoup4",
    "lxml": "lxml",
    "pandas": "pandas",
    "numpy": "numpy",
    "matplotlib": "matplotlib",
    "seaborn": "seaborn",
    "plotly": "plotly",
    "requests": "requests",
    "openpyxl": "openpyxl",
    "docx": "python-docx",
    "reportlab": "reportlab",
    "xlrd": "xlrd",
    "chardet": "chardet",
    "tabulate": "tabulate",
    "jinja2": "Jinja2",
    "dateutil": "python-dateutil",
    "pytz": "pytz",
    "tqdm": "tqdm",
    "colorama": "colorama",
    "rich": "rich",
    "click": "click",
}

# Dangerous modules/functions that are BLOCKED
BLOCKED_MODULES = {
    "os", "sys", "subprocess", "shutil", "socket", "ctypes",
    "importlib", "builtins", "code", "codeop", "compile",
    "compileall", "py_compile", "zipimport", "pkgutil",
    "signal", "multiprocessing", "threading", "_thread",
    "webbrowser", "http", "urllib", "xmlrpc", "smtplib",
    "ftplib", "telnetlib", "poplib", "imaplib", "nntplib",
    "ssl", "tempfile", "glob", "fnmatch", "shlex",
}

BLOCKED_CALLS = {
    "exec", "eval", "compile", "open", "__import__",
    "globals", "locals", "dir", "vars", "getattr",
    "setattr", "delattr", "breakpoint", "exit", "quit",
    "input",
}

# Standard library modules that don't need pip install
STDLIB = {
    "abc", "argparse", "array", "ast", "asyncio", "base64",
    "bisect", "calendar", "cmath", "collections", "colorsys",
    "configparser", "contextlib", "copy", "csv", "dataclasses",
    "datetime", "decimal", "difflib", "enum", "fcntl",
    "filecmp", "fractions", "functools", "gc", "getpass",
    "gzip", "hashlib", "heapq", "html", "inspect",
    "io", "itertools", "json", "keyword", "linecache",
    "locale", "logging", "math", "mimetypes", "numbers",
    "operator", "pathlib", "pdb", "pickle", "plistlib",
    "pprint", "profile", "queue", "random", "re",
    "secrets", "shelve", "sqlite3", "statistics", "string",
    "struct", "textwrap", "time", "timeit", "token",
    "tokenize", "trace", "traceback", "turtle", "types",
    "typing", "unicodedata", "unittest", "urllib", "uuid",
    "warnings", "weakref", "xml", "zipfile", "zlib",
}


# ── Models ───────────────────────────────────────────────────────────────

class GenerateRequest(BaseModel):
    code: str


# ── Security Scanner ─────────────────────────────────────────────────────

def scan_code(code: str) -> tuple[bool, str, list[str]]:
    """
    Scan code for dangerous operations and extract import names.
    Returns (is_safe, error_message, list_of_imports).
    """
    try:
        tree = ast.parse(code)
    except SyntaxError as e:
        return False, f"Syntax error: {e}", []

    imports = set()

    for node in ast.walk(tree):
        # Check imports
        if isinstance(node, ast.Import):
            for alias in node.names:
                top = alias.name.split(".")[0]
                if top in BLOCKED_MODULES:
                    return False, f"Blocked module: {top}", []
                imports.add(top)

        elif isinstance(node, ast.ImportFrom):
            if node.module:
                top = node.module.split(".")[0]
                if top in BLOCKED_MODULES:
                    return False, f"Blocked module: {top}", []
                imports.add(top)

        # Check function calls
        elif isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                if node.func.id in BLOCKED_CALLS:
                    return False, f"Blocked function: {node.func.id}()", []
            elif isinstance(node.func, ast.Attribute):
                if node.func.attr in {"system", "popen", "exec", "rmtree", "remove", "unlink"}:
                    return False, f"Blocked call: .{node.func.attr}()", []

    return True, "", list(imports)


def get_pip_packages(imports: list[str]) -> list[str]:
    """Convert import names to pip package names."""
    packages = []
    for imp in imports:
        if imp in STDLIB:
            continue
        pkg = ALLOWED_PACKAGES.get(imp, imp)
        packages.append(pkg)
    return list(set(packages))


# ── Environment Management ───────────────────────────────────────────────

ENV_CACHE_DIR = Path("/tmp/smartppt_envs")
ENV_CACHE_DIR.mkdir(parents=True, exist_ok=True)

JOBS_DIR = Path("/tmp/smartppt_jobs")
JOBS_DIR.mkdir(parents=True, exist_ok=True)


def get_or_create_venv(packages: list[str]) -> Path:
    """Get or create a cached virtual environment for the given packages."""
    # Hash the package list for caching
    pkg_key = hashlib.sha256("|".join(sorted(packages)).encode()).hexdigest()[:16]
    venv_dir = ENV_CACHE_DIR / f"env_{pkg_key}"

    if venv_dir.exists():
        return venv_dir

    # Create venv
    subprocess.run(
        [sys.executable, "-m", "venv", str(venv_dir)],
        check=True,
        capture_output=True,
        timeout=60,
    )

    # Install packages
    if packages:
        pip_exe = venv_dir / "bin" / "pip"
        if not pip_exe.exists():
            pip_exe = venv_dir / "Scripts" / "pip.exe"

        subprocess.run(
            [str(pip_exe), "install", "--quiet", "--no-warn-script-location"] + packages,
            check=True,
            capture_output=True,
            timeout=300,
        )

    return venv_dir


# ── Endpoints ────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {"status": "ok", "message": "SmartPPT Backend is running"}


@app.get("/")
async def root():
    return {
        "app": "SmartPPT Builder API",
        "version": "1.0.0",
        "docs": "/docs",
    }


@app.post("/api/generate")
async def generate(
    request: GenerateRequest,
    authorization: Optional[str] = Header(None),
):
    """Generate a .pptx file from Python code."""

    code = request.code.strip()
    if not code:
        raise HTTPException(status_code=400, detail="No code provided")

    # 0. Auth Verification
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
        
    id_token = authorization.split("Bearer ")[1]
    try:
        # Verify token securely to prevent API abuse
        decoded_token = firebase_auth.verify_id_token(id_token)
        user_id = decoded_token['uid']
    except Exception as e:
        raise HTTPException(status_code=401, detail=f"Unauthorized: Invalid token ({str(e)})")

    if len(code) > 50000:
        raise HTTPException(status_code=400, detail="Code too long (max 50,000 chars)")

    # 1. Security scan
    is_safe, error, imports = scan_code(code)
    if not is_safe:
        raise HTTPException(status_code=400, detail=f"Security check failed: {error}")

    # 2. Determine packages
    packages = get_pip_packages(imports)

    # Ensure python-pptx is always available
    if "python-pptx" not in packages:
        packages.append("python-pptx")

    # 3. Get/create cached environment
    try:
        venv_dir = get_or_create_venv(packages)
    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=504, detail="Environment setup timed out")
    except subprocess.CalledProcessError as e:
        raise HTTPException(status_code=500, detail=f"Failed to create environment: {e.stderr}")

    # 4. Create job directory
    job_dir = JOBS_DIR / hashlib.sha256(code.encode()).hexdigest()[:16]
    if job_dir.exists():
        shutil.rmtree(job_dir)
    job_dir.mkdir(parents=True)

    script_path = job_dir / "script.py"
    script_path.write_text(code, encoding="utf-8")

    # 5. Run the script
    python_exe = venv_dir / "bin" / "python"
    if not python_exe.exists():
        python_exe = venv_dir / "Scripts" / "python.exe"

    try:
        result = subprocess.run(
            [str(python_exe), str(script_path)],
            cwd=str(job_dir),
            capture_output=True,
            text=True,
            timeout=120,
        )
    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=504, detail="Script execution timed out (120s limit)")

    if result.returncode != 0:
        err = result.stderr.strip() or "(no error output)"
        raise HTTPException(
            status_code=400,
            detail=f"Script failed (exit code {result.returncode}):\n{err}",
        )

    # 6. Find generated .pptx files
    pptx_files = list(job_dir.glob("*.pptx"))
    if not pptx_files:
        raise HTTPException(
            status_code=400,
            detail="Script ran successfully but no .pptx file was produced. "
                   "Make sure your code calls prs.save('filename.pptx')",
        )

    # Return the last created .pptx
    pptx_file = sorted(pptx_files, key=lambda p: p.stat().st_mtime)[-1]

    return FileResponse(
        path=str(pptx_file),
        filename=pptx_file.name,
        media_type="application/vnd.openxmlformats-officedocument.presentationml.presentation",
    )


# ── Cleanup old jobs periodically ────────────────────────────────────────

@app.on_event("startup")
async def cleanup_old_jobs():
    """Clean up job directories older than 1 hour."""
    import time
    cutoff = time.time() - 3600
    if JOBS_DIR.exists():
        for job in JOBS_DIR.iterdir():
            if job.is_dir() and job.stat().st_mtime < cutoff:
                shutil.rmtree(job, ignore_errors=True)