from __future__ import annotations
import base64, io, os, sqlite3
from pathlib import Path
import requests, stripe
from dotenv import load_dotenv
from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from PIL import Image

load_dotenv()
APP_NAME=os.getenv("APP_NAME","Historic Photo Colorizer")
RUNPOD_API_KEY=os.getenv("RUNPOD_API_KEY","")
RUNPOD_ENDPOINT_ID=os.getenv("RUNPOD_ENDPOINT_ID","")
REQUIRE_PAYMENT=os.getenv("REQUIRE_PAYMENT","false").lower()=="true"
STRIPE_SECRET_KEY=os.getenv("STRIPE_SECRET_KEY","")
STRIPE_PRICE_ID=os.getenv("STRIPE_PRICE_ID","")
PUBLIC_BASE_URL=os.getenv("PUBLIC_BASE_URL","http://localhost:8000").rstrip("/")
MAX_UPLOAD_MB=int(os.getenv("MAX_UPLOAD_MB","6"))
stripe.api_key=STRIPE_SECRET_KEY

BASE=Path(__file__).resolve().parent
DB=BASE/"payments.sqlite3"
app=FastAPI(title=APP_NAME)
app.mount("/static",StaticFiles(directory=BASE/"static"),name="static")
templates=Jinja2Templates(directory=BASE/"templates")

def init_db():
    with sqlite3.connect(DB) as con:
        con.execute("CREATE TABLE IF NOT EXISTS used_sessions(session_id TEXT PRIMARY KEY,used_at DATETIME DEFAULT CURRENT_TIMESTAMP)")
init_db()

def session_used(sid):
    with sqlite3.connect(DB) as con:
        return con.execute("SELECT 1 FROM used_sessions WHERE session_id=?",(sid,)).fetchone() is not None

def mark_used(sid):
    with sqlite3.connect(DB) as con:
        con.execute("INSERT OR IGNORE INTO used_sessions(session_id) VALUES (?)",(sid,))
        con.commit()

def verify_paid(sid):
    if not REQUIRE_PAYMENT: return True
    if not sid or not STRIPE_SECRET_KEY or session_used(sid): return False
    try:
        s=stripe.checkout.Session.retrieve(sid)
        return s.payment_status=="paid"
    except Exception:
        return False

@app.get("/",response_class=HTMLResponse)
def home(request:Request,session_id:str=""):
    paid=verify_paid(session_id) if session_id else (not REQUIRE_PAYMENT)
    return templates.TemplateResponse("index.html",{"request":request,"app_name":APP_NAME,"require_payment":REQUIRE_PAYMENT,"paid":paid,"session_id":session_id,"max_upload_mb":MAX_UPLOAD_MB})

@app.post("/create-checkout")
def create_checkout():
    if not REQUIRE_PAYMENT: return RedirectResponse("/",status_code=303)
    if not STRIPE_SECRET_KEY or not STRIPE_PRICE_ID: raise HTTPException(500,"Stripe is not configured")
    s=stripe.checkout.Session.create(
        mode="payment",
        line_items=[{"price":STRIPE_PRICE_ID,"quantity":1}],
        success_url=f"{PUBLIC_BASE_URL}/?session_id={{CHECKOUT_SESSION_ID}}",
        cancel_url=f"{PUBLIC_BASE_URL}/",
    )
    return RedirectResponse(s.url,status_code=303)

@app.post("/api/colorize")
async def colorize(
    image:UploadFile=File(...),
    model:str=Form("artistic"),
    quality:str=Form("balanced"),
    strength:float=Form(.80),
    filter_style:str=Form("historical"),
    session_id:str=Form("")
):
    if REQUIRE_PAYMENT and not verify_paid(session_id):
        raise HTTPException(402,"A valid unused payment is required")
    if not RUNPOD_API_KEY or not RUNPOD_ENDPOINT_ID:
        raise HTTPException(500,"RunPod is not configured")

    raw=await image.read()
    if len(raw)>MAX_UPLOAD_MB*1024*1024:
        raise HTTPException(413,f"Image exceeds {MAX_UPLOAD_MB} MB")
    try:
        p=Image.open(io.BytesIO(raw)); p.verify()
    except Exception:
        raise HTTPException(400,"Uploaded file is not a valid image")

    payload={"input":{
        "image_base64":base64.b64encode(raw).decode("ascii"),
        "model":model,"quality":quality,"strength":strength,"filter":filter_style
    }}
    url=f"https://api.runpod.ai/v2/{RUNPOD_ENDPOINT_ID}/runsync?wait=300000"
    r=requests.post(url,headers={"Authorization":f"Bearer {RUNPOD_API_KEY}","Content-Type":"application/json"},json=payload,timeout=320)
    if not r.ok: raise HTTPException(502,f"GPU service error: {r.text[:500]}")
    data=r.json()
    if data.get("status")!="COMPLETED":
        raise HTTPException(502,f"GPU job did not complete: {data.get('status')}")
    out=data.get("output") or {}
    if not out.get("image_base64"): raise HTTPException(502,"GPU service returned no image")
    if REQUIRE_PAYMENT: mark_used(session_id)
    return JSONResponse({"image_base64":out["image_base64"],"mime_type":out.get("mime_type","image/jpeg"),"meta":{k:out.get(k) for k in ["model","quality","filter","render_factor","width","height"]}})

@app.get("/health")
def health():
    return {"ok":True,"app":APP_NAME}