main.py

import os
import json
from datetime import datetime
from io import BytesIO
from typing import List, Dict, Any, Optional

import numpy as np
import requests
from dotenv import load_dotenv
from fastapi import FastAPI, Request, UploadFile, File, HTTPException
from fastapi.responses import RedirectResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.status import HTTP_303_SEE_OTHER
from PIL import Image

# ----------------------------------------------------
# Initial setup
# ----------------------------------------------------
load_dotenv()

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ENC_FOLDER = os.path.join(BASE_DIR, "encrypted_images")
META_FILE = os.path.join(BASE_DIR, "image_metadata.json")

os.makedirs(ENC_FOLDER, exist_ok=True)

PINATA_JWT = os.getenv("PINATA_JWT")
PINATA_GATEWAY_URL = os.getenv("PINATA_GATEWAY_URL", "https://gateway.pinata.cloud/ipfs")
PINATA_UPLOAD_URL = "https://api.pinata.cloud/pinning/pinFileToIPFS"

app = FastAPI(title="Visual Crypto + IPFS Demo")

app.mount("/encrypted_images", StaticFiles(directory=ENC_FOLDER), name="encrypted_images")
templates = Jinja2Templates(directory=os.path.join(BASE_DIR, "templates"))


# ----------------------------------------------------
# Helpers: metadata
# ----------------------------------------------------
def load_metadata() -> List[Dict[str, Any]]:
    if not os.path.exists(META_FILE):
        return []
    try:
        with open(META_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return []


def save_metadata(data: List[Dict[str, Any]]) -> None:
    with open(META_FILE, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def next_id(data: List[Dict[str, Any]]) -> int:
    if not data:
        return 1
    return max(item["id"] for item in data) + 1


# ----------------------------------------------------
# Helpers: visual encryption (image XOR)
# ----------------------------------------------------
def pil_image_from_upload(upload: UploadFile) -> Image.Image:
    contents = upload.file.read()
    if not contents:
        raise HTTPException(status_code=400, detail="Empty file upload")
    try:
        img = Image.open(BytesIO(contents)).convert("RGBA")
        return img
    except Exception:
        raise HTTPException(status_code=400, detail="Unsupported or invalid image file")


def generate_noise_key_and_encrypt(original: Image.Image):
    """
    original: PIL RGBA image
    returns: (encrypted_image, key_image) both PIL RGBA
    """
    arr = np.array(original, dtype=np.uint8)
    # same shape random noise
    key_arr = np.random.randint(0, 256, size=arr.shape, dtype=np.uint8)
    enc_arr = np.bitwise_xor(arr, key_arr)

    key_img = Image.fromarray(key_arr, mode="RGBA")
    enc_img = Image.fromarray(enc_arr, mode="RGBA")
    return enc_img, key_img


def decrypt_with_key(enc_img: Image.Image, key_img: Image.Image) -> Image.Image:
    enc_arr = np.array(enc_img.convert("RGBA"), dtype=np.uint8)
    key_arr = np.array(key_img.convert("RGBA"), dtype=np.uint8)

    if enc_arr.shape != key_arr.shape:
        raise HTTPException(status_code=500, detail="Encrypted and key image shapes do not match")

    orig_arr = np.bitwise_xor(enc_arr, key_arr)
    return Image.fromarray(orig_arr, mode="RGBA")


# ----------------------------------------------------
# Helpers: Pinata
# ----------------------------------------------------
def ensure_pinata_configured():
    if not PINATA_JWT:
        raise HTTPException(
            status_code=500,
            detail="PINATA_JWT is not configured. Set it in .env or environment variables.",
        )


def upload_key_to_pinata(key_img: Image.Image, filename: str) -> str:
    ensure_pinata_configured()

    # Save key image into memory as PNG
    buf = BytesIO()
    key_img.save(buf, format="PNG")
    buf.seek(0)

    headers = {
        "Authorization": PINATA_JWT,
    }
    files = {
        "file": (filename, buf, "image/png"),
    }

    resp = requests.post(PINATA_UPLOAD_URL, headers=headers, files=files, timeout=60)
    if resp.status_code != 200:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to upload key to Pinata: {resp.text}",
        )

    data = resp.json()
    cid = data.get("IpfsHash")
    if not cid:
        raise HTTPException(
            status_code=500,
            detail=f"Pinata response missing IpfsHash: {data}",
        )
    return cid


def fetch_key_from_ipfs(cid: str) -> Image.Image:
    url = f"{PINATA_GATEWAY_URL.rstrip('/')}/{cid}"
    resp = requests.get(url, timeout=60)
    if resp.status_code != 200:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to fetch key from IPFS gateway: {resp.status_code}",
        )
    try:
        key_img = Image.open(BytesIO(resp.content)).convert("RGBA")
        return key_img
    except Exception:
        raise HTTPException(
            status_code=500,
            detail="Failed to decode key image from IPFS content",
        )


# ----------------------------------------------------
# Routes
# ----------------------------------------------------
@app.get("/")
async def index(request: Request):
    items = load_metadata()
    has_pinata = PINATA_JWT is not None
    return templates.TemplateResponse(
        "index.html",
        {
            "request": request,
            "items": items,
            "has_pinata": has_pinata,
        },
    )


@app.post("/upload_original")
async def upload_original(request: Request, file: UploadFile = File(...)):
    """
    1. Read original image
    2. Generate random key image (noise)
    3. Encrypt original using XOR → encrypted noise image
    4. Store encrypted image locally
    5. Upload key image to IPFS via Pinata
    6. Save metadata (enc_filename, CID, timestamps)
    """
    if not file:
        raise HTTPException(status_code=400, detail="No file uploaded")

    original_img = pil_image_from_upload(file)
    enc_img, key_img = generate_noise_key_and_encrypt(original_img)

    # Generate filenames
    now = datetime.utcnow().strftime("%Y%m%d%H%M%S")
    base_name = os.path.splitext(file.filename or "image")[0]
    enc_filename = f"enc_{now}_{base_name}.png"
    key_filename = f"key_{now}_{base_name}.png"

    # Save encrypted image locally
    enc_path = os.path.join(ENC_FOLDER, enc_filename)
    enc_img.save(enc_path, format="PNG")

    # Upload key to IPFS
    cid = upload_key_to_pinata(key_img, key_filename)

    # Update metadata
    data = load_metadata()
    item_id = next_id(data)
    data.append(
        {
            "id": item_id,
            "enc_filename": enc_filename,
            "key_cid": cid,
            "original_name": file.filename,
            "created_at": datetime.utcnow().isoformat() + "Z",
            "last_decrypted_at": None,
        }
    )
    save_metadata(data)

    # Redirect back to index
    return RedirectResponse(url="/", status_code=HTTP_303_SEE_OTHER)


@app.get("/decrypt/{item_id}")
async def decrypt_item(item_id: int):
    """
    Fetch key image from IPFS using CID, XOR with encrypted image, and
    stream the recovered original image back as PNG.
    """
    data = load_metadata()
    item: Optional[Dict[str, Any]] = next((x for x in data if x["id"] == item_id), None)
    if not item:
        raise HTTPException(status_code=404, detail="Item not found")

    enc_path = os.path.join(ENC_FOLDER, item["enc_filename"])
    if not os.path.exists(enc_path):
        raise HTTPException(status_code=404, detail="Encrypted image file not found on server")

    enc_img = Image.open(enc_path).convert("RGBA")
    key_img = fetch_key_from_ipfs(item["key_cid"])
    orig_img = decrypt_with_key(enc_img, key_img)

    # Update metadata last_decrypted_at
    for x in data:
        if x["id"] == item_id:
            x["last_decrypted_at"] = datetime.utcnow().isoformat() + "Z"
            break
    save_metadata(data)

    buf = BytesIO()
    orig_img.save(buf, format="PNG")
    buf.seek(0)
    return StreamingResponse(buf, media_type="image/png")


# Simple healthcheck
@app.get("/health")
async def health():
    return {"status": "ok"}