import os, sqlite3, datetime, secrets, hashlib, subprocess, json, time, functools
from fastapi import FastAPI, Request, Depends, HTTPException, Form
from fastapi.responses import HTMLResponse, FileResponse, RedirectResponse, JSONResponse, PlainTextResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from jose import jwt
from jinja2 import Environment, FileSystemLoader, select_autoescape
from config import *
from models import get_db, init_db, hash_password, verify_password
from auth import create_access_token, get_current_user
from signing import build_leaf_cert
from cryptography.hazmat.primitives import serialization

app = FastAPI(title="CertAuth Key Vault")
app.mount("/static", StaticFiles(directory="/opt/certauth/api/static"), name="static")

_login_attempts = {}
_LOGIN_MAX_ATTEMPTS = 5
_LOGIN_WINDOW_SECONDS = 900

_csrf_secret = secrets.token_hex(32)


def _check_rate_limit(client_ip: str) -> bool:
    now = time.time()
    if client_ip not in _login_attempts:
        _login_attempts[client_ip] = []
    _login_attempts[client_ip] = [
        t for t in _login_attempts[client_ip] if now - t < _LOGIN_WINDOW_SECONDS
    ]
    if len(_login_attempts[client_ip]) >= _LOGIN_MAX_ATTEMPTS:
        return False
    _login_attempts[client_ip].append(now)
    return True


def _generate_csrf_token() -> str:
    return secrets.token_hex(32)


def _verify_csrf_token(request: Request, token: str) -> bool:
    stored = request.cookies.get("csrf_token")
    if not stored or not token:
        return False
    return secrets.compare_digest(stored, token)


def _set_csrf_cookie(resp):
    token = _generate_csrf_token()
    resp.set_cookie("csrf_token", token, httponly=False, samesite="strict", path="/")
    return token


jinja_env = Environment(
    loader=FileSystemLoader("/opt/certauth/api/templates"),
    autoescape=select_autoescape(["html"])
)

@app.on_event("startup")
def startup():
    init_db()
    try:
        with open(ROOT_CA_PATH) as f: root = f.read()
        with open(INT_CA_PATH) as f: inter = f.read()
        with open(CA_CHAIN_PATH, "w") as f: f.write(inter + "\n" + root)
    except: pass

def render(name, ctx):
    return HTMLResponse(jinja_env.get_template(name).render(**ctx))

class LoginRequest(BaseModel):
    username: str
    password: str

@app.post("/api/token")
async def login(req: LoginRequest, request: Request):
    client_ip = request.client.host
    if not _check_rate_limit(client_ip):
        raise HTTPException(429, "Too many login attempts. Try again later.")
    conn = get_db()
    row = conn.execute("SELECT * FROM users WHERE username = ?", (req.username,)).fetchone()
    conn.close()
    if not row or not verify_password(req.password, row["password_hash"]):
        raise HTTPException(401, "Invalid credentials")
    token = create_access_token({"sub": req.username})
    return {"access_token": token, "token_type": "bearer"}

@app.get("/api/me")
async def me(user: str = Depends(get_current_user)):
    return {"username": user}

@app.get("/api/domains")
async def list_domains(user: str = Depends(get_current_user)):
    conn = get_db()
    rows = conn.execute("SELECT * FROM domains ORDER BY created_at DESC").fetchall()
    conn.close()
    return [dict(r) for r in rows]

@app.post("/api/domains")
async def create_domain(name: str = Form(...), description: str = Form(""),
                        user: str = Depends(get_current_user)):
    conn = get_db()
    cur = conn.cursor()
    cur.execute("INSERT INTO domains (name, description, created_by) VALUES (?,?,?)",
                 (name, description, 1))
    conn.commit()
    conn.close()
    return {"status": "ok"}

@app.get("/api/certs")
async def list_certs(user: str = Depends(get_current_user)):
    conn = get_db()
    rows = conn.execute("SELECT c.*, d.name as domain_name FROM certificates c LEFT JOIN domains d ON c.domain_id = d.id ORDER BY c.created_at DESC").fetchall()
    conn.close()
    return [dict(r) for r in rows]

@app.post("/api/certs/request")
async def request_cert(cn: str = Form(...), sans: str = Form(""),
                       days: int = Form(365), domain_id: int = Form(0),
                       user: str = Depends(get_current_user)):
    conn = get_db()
    cur = conn.cursor()
    cur.execute("INSERT INTO certificates (domain_id, subject, san, status, created_by) VALUES (?,?,?,?,?)",
                 (domain_id, cn, sans, "pending", 1))
    conn.commit()
    cid = cur.lastrowid
    conn.close()
    return {"status": "ok", "id": cid}

@app.post("/api/certs/{cert_id}/sign")
async def sign_cert(cert_id: int, user: str = Depends(get_current_user)):
    conn = get_db()
    row = conn.execute("SELECT * FROM certificates WHERE id = ?", (cert_id,)).fetchone()
    if not row or row["status"] != "pending":
        conn.close()
        raise HTTPException(400, "Not found or already signed")
    conn.close()
    result, err = build_leaf_cert(row["subject"], row["san"], 365)
    if err:
        import html as h
        raise HTTPException(500, f"Signing failed: {h.escape(str(err))}")
    cf = f"/etc/ssl/ca/issued/cert-{result['serial']}.crt"
    kf = f"/etc/ssl/ca/issued/cert-{result['serial']}.key"
    open(cf, "w").write(result["cert_pem"])
    open(kf, "w").write(result["key_pem"])
    os.chmod(cf, 0o640); os.chmod(kf, 0o600)
    conn = get_db()
    conn.execute("UPDATE certificates SET status=?, serial=?, cert_path=?, issued_at=?, expires_at=? WHERE id=?",
                 ("issued", result["serial"], cf, datetime.datetime.now().isoformat(), result["expires_at"], cert_id))
    conn.commit(); conn.close()
    return {"status": "ok", "serial": result["serial"]}

@app.get("/api/certs/{cert_id}/pem")
async def download_pem(cert_id: int, request: Request = None):
    """Download cert + chain as bundled PEM."""
    user = get_user_from_cookie(request)
    if not user: raise HTTPException(401, "Login required")
    conn = get_db()
    row = conn.execute("SELECT * FROM certificates WHERE id = ?", (cert_id,)).fetchone()
    conn.close()
    if not row or row["status"] != "issued": raise HTTPException(404)
    pem_path = f"/var/lib/certauth/tmp/cert-{row['serial']}.pem"
    with open(row["cert_path"]) as f:
        cert_pem = f.read()
    with open(CA_CHAIN_PATH) as f:
        chain_pem = f.read()
    with open(pem_path, "w") as f:
        f.write(cert_pem.rstrip() + "\n" + chain_pem)
    return FileResponse(pem_path, media_type="application/x-pem-file", filename=f"cert-{row['serial']}.pem")

@app.get("/api/certs/{cert_id}/pfx")
async def download_pfx(cert_id: int, password: str = "certauth", request: Request = None):
    """Download cert + key + chain as PKCS12/PFX."""
    user = get_user_from_cookie(request)
    if not user: raise HTTPException(401, "Login required")
    conn = get_db()
    row = conn.execute("SELECT * FROM certificates WHERE id = ?", (cert_id,)).fetchone()
    conn.close()
    if not row or row["status"] != "issued": raise HTTPException(404)
    kf = row["cert_path"].replace(".crt", ".key")
    if not os.path.exists(kf): raise HTTPException(404)
    from cryptography.hazmat.primitives.serialization import pkcs12, BestAvailableEncryption
    from cryptography import x509
    with open(row["cert_path"], "rb") as f:
        leaf = x509.load_pem_x509_certificate(f.read())
    with open(kf, "rb") as f:
        key = serialization.load_pem_private_key(f.read(), password=None)
    chain_certs = []
    with open(CA_CHAIN_PATH, "rb") as f:
        for cert_pem in f.read().split(b"-----END CERTIFICATE-----"):
            cert_pem = cert_pem.strip()
            if cert_pem:
                chain_certs.append(x509.load_pem_x509_certificate(cert_pem + b"\n-----END CERTIFICATE-----"))
    pfx_data = pkcs12.serialize_key_and_certificates(
        name=row["subject"].encode(),
        key=key,
        cert=leaf,
        cas=chain_certs or None,
        encryption_algorithm=BestAvailableEncryption(password.encode())
    )
    pfx_path = f"/var/lib/certauth/tmp/cert-{row['serial']}.pfx"
    with open(pfx_path, "wb") as f:
        f.write(pfx_data)
    return FileResponse(pfx_path, media_type="application/x-pkcs12", filename=f"cert-{row['serial']}.pfx")

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

@app.get("/api/ca-chain")
async def ca_chain(): return FileResponse(CA_CHAIN_PATH, filename="ca-chain.crt")

def get_user_from_cookie(request: Request):
    token = request.cookies.get("token")
    if not token:
        return None
    try:
        return jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
    except:
        return None

@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    user = get_user_from_cookie(request)
    if not user:
        return RedirectResponse("/login", status_code=302)
    conn = get_db()
    certs = conn.execute("SELECT c.*, d.name as domain_name FROM certificates c LEFT JOIN domains d ON c.domain_id = d.id ORDER BY c.created_at DESC LIMIT 20").fetchall()
    domains = conn.execute("SELECT * FROM domains").fetchall()
    p = conn.execute("SELECT COUNT(*) as c FROM certificates WHERE status = ?", ("pending",)).fetchone()["c"]
    i = conn.execute("SELECT COUNT(*) as c FROM certificates WHERE status = ?", ("issued",)).fetchone()["c"]
    conn.close()
    return render("dashboard.html", {"request": request, "user": user,
        "certs": [dict(r) for r in certs], "domains": [dict(r) for r in domains],
        "pending": p, "issued": i})

@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
    return render("login.html", {"request": request, "error": None})

@app.post("/login")
async def login_post(username: str = Form(...), password: str = Form(...),
                     request: Request = None):
    client_ip = request.client.host if request else "unknown"
    if not _check_rate_limit(client_ip):
        return render("login.html", {"request": None, "error": "Too many login attempts. Try again later."})
    conn = get_db()
    row = conn.execute("SELECT * FROM users WHERE username = ?", (username,)).fetchone()
    conn.close()
    if not row or not verify_password(password, row["password_hash"]):
        return render("login.html", {"request": None, "error": "Invalid credentials"})
    token = create_access_token({"sub": username})
    resp = RedirectResponse("/", status_code=302)
    resp.set_cookie("token", token, httponly=True, samesite="strict", secure=True, path="/")
    _set_csrf_cookie(resp)
    return resp



@app.post("/api/certs/{cert_id}/sign/web")
async def sign_cert_web(cert_id: int, csrf_token: str = Form(""), request: Request = None):
    user = get_user_from_cookie(request)
    if not user:
        return RedirectResponse("/login", status_code=302)
    if not _verify_csrf_token(request, csrf_token):
        raise HTTPException(403, "Invalid CSRF token")
    conn = get_db()
    row = conn.execute("SELECT * FROM certificates WHERE id = ?", (cert_id,)).fetchone()
    if not row or row["status"] != "pending":
        conn.close()
        raise HTTPException(400, "Not found or already issued")
    conn.close()
    import html as html_lib
    try:
        result, err = build_leaf_cert(row["subject"], row["san"], 365)
        if err:
            safe_err = html_lib.escape(str(err))
            return HTMLResponse(f'<span class="text-red-400">Issue failed: {safe_err}</span>')
        cf = f"/etc/ssl/ca/issued/cert-{result['serial']}.crt"
        kf = f"/etc/ssl/ca/issued/cert-{result['serial']}.key"
        open(cf, "w").write(result["cert_pem"])
        open(kf, "w").write(result["key_pem"])
        os.chmod(cf, 0o640)
        os.chmod(kf, 0o600)
        conn2 = get_db()
        conn2.execute("UPDATE certificates SET status=?, serial=?, cert_path=?, issued_at=?, expires_at=? WHERE id=?",
                      ("issued", result["serial"], cf, datetime.datetime.now().isoformat(), result["expires_at"], cert_id))
        conn2.commit()
        conn2.close()
        return HTMLResponse(f'<span class="text-green-400">Issued! <a href="/api/certs/{cert_id}/pem" class="underline">PEM</a> | <a href="/api/certs/{cert_id}/pfx" class="underline">PFX</a> | <a href="/certs" class="underline">Refresh</a></span>')
    except Exception as ex:
        safe_ex = html_lib.escape(str(ex))
        return HTMLResponse(f'<span class="text-red-400">Issue failed: {safe_ex}</span>')


@app.get("/logout")
async def logout():
    resp = RedirectResponse("/login", status_code=302)
    resp.delete_cookie("token", path="/")
    return resp

# --- Web API (cookie auth) ---
@app.post("/api/domains/web")
async def create_domain_web(name: str = Form(...), description: str = Form(""),
                            csrf_token: str = Form(""), request: Request = None):
    user = get_user_from_cookie(request)
    if not user:
        return RedirectResponse("/login", status_code=302)
    if not _verify_csrf_token(request, csrf_token):
        raise HTTPException(403, "Invalid CSRF token")
    conn = get_db()
    cur = conn.cursor()
    cur.execute("INSERT INTO domains (name, description, created_by) VALUES (?,?,?)", (name, description, 1))
    conn.commit()
    conn.close()
    return HTMLResponse('<span class="text-green-400">Domain registered! <a href="/domains" class="underline">Refresh</a></span>')

@app.post("/api/certs/web/request")
async def request_cert_web(cn: str = Form(...), sans: str = Form(""), days: int = Form(365),
                           domain_id: int = Form(0), csrf_token: str = Form(""),
                           request: Request = None):
    user = get_user_from_cookie(request)
    if not user:
        return RedirectResponse("/login", status_code=302)
    if not _verify_csrf_token(request, csrf_token):
        raise HTTPException(403, "Invalid CSRF token")
    conn = get_db()
    cur = conn.cursor()
    if domain_id == 0:
        cur.execute("SELECT id FROM domains WHERE name=?", (cn,))
        row = cur.fetchone()
        domain_id = row[0] if row else None
    cur.execute("INSERT INTO certificates (domain_id, subject, san, status, created_by) VALUES (?,?,?,?,?)", (domain_id, cn, sans, "pending", 1))
    conn.commit()
    conn.close()
    return HTMLResponse('<span class="text-green-400">Certificate requested! Click Issue below. <a href="/certs" class="underline">Refresh</a></span>')

@app.get("/domains", response_class=HTMLResponse)
async def domains_page(request: Request):
    user = get_user_from_cookie(request)
    if not user: return RedirectResponse("/login", status_code=302)
    conn = get_db()
    rows = conn.execute("SELECT * FROM domains ORDER BY created_at DESC").fetchall()
    conn.close()
    return render("domains.html", {"request": request, "user": user, "domains": [dict(r) for r in rows]})

@app.get("/certs", response_class=HTMLResponse)
async def certs_page(request: Request):
    user = get_user_from_cookie(request)
    if not user: return RedirectResponse("/login", status_code=302)
    conn = get_db()
    rows = conn.execute("SELECT c.*, d.name as domain_name FROM certificates c LEFT JOIN domains d ON c.domain_id = d.id ORDER BY c.created_at DESC").fetchall()
    domains = conn.execute("SELECT * FROM domains").fetchall()
    conn.close()
    return render("certs.html", {"request": request, "user": user, "certs": [dict(r) for r in rows],
                                  "domains": [dict(r) for r in domains]})

@app.get("/history", response_class=HTMLResponse)
async def history_page(request: Request):
    user = get_user_from_cookie(request)
    if not user: return RedirectResponse("/login", status_code=302)
    conn = get_db()
    rows = conn.execute("SELECT c.*, d.name as domain_name FROM certificates c LEFT JOIN domains d ON c.domain_id = d.id ORDER BY c.created_at DESC").fetchall()
    conn.close()
    return render("history.html", {"request": request, "user": user, "certs": [dict(r) for r in rows]})

@app.get("/setup", response_class=HTMLResponse)
async def setup_page(request: Request):
    user = get_user_from_cookie(request)
    if not user: return RedirectResponse("/login", status_code=302)
    return render("setup.html", {"request": request, "user": user})

@app.get("/setup.sh")
async def setup_sh():
    """One-liner bash setup script for Linux/macOS."""
    script = r'''#!/bin/bash
set -e
# CertAuth CA Chain Installer
# Usage: curl -sL http://<certauth-ip>/setup.sh | bash
#        curl -sL http://<certauth-ip>/setup.sh | sudo bash

# Auto-detect CertAuth server IP
DETECTED_IP=""
if [[ -n "$1" ]]; then
    DETECTED_IP="$1"
elif [[ -n "$CERTAUTH_IP" ]]; then
    DETECTED_IP="$CERTAUTH_IP"
else
    # Try Linux hostname -I first
    DETECTED_IP=$(hostname -I 2>/dev/null | awk '{print $1}') || true
    # Fallback: ip route (Linux)
    [[ -z "$DETECTED_IP" ]] && DETECTED_IP=$(ip route get 1 2>/dev/null | awk '{print $7}' | head -1) || true
    # Fallback: ifconfig (macOS/BSD)
    [[ -z "$DETECTED_IP" ]] && DETECTED_IP=$(ifconfig 2>/dev/null | grep -E '^\s+(inet )' | awk '{print $2}' | grep -v '127.0.0.1' | head -1) || true
    # Fallback: networksetup (macOS only)
    [[ -z "$DETECTED_IP" ]] && DETECTED_IP=$(networksetup -getinfo $(networksetup -listallhardwareports 2>/dev/null | awk '/Hardware Port:/ {getline; gsub(/^[ \t]+/, ""); print}') 2>/dev/null | grep 'IP address:' | awk '{print $3}') || true
fi

# Prompt if auto-detection failed
if [[ -z "$DETECTED_IP" ]]; then
    read -r -p "Enter CertAuth server IP (e.g., 192.168.8.248): " DETECTED_IP
fi

CHAIN_URL="http://$DETECTED_IP/api/ca-chain"

echo "Downloading CA chain..."
curl -sLk "$CHAIN_URL" -o /tmp/ca-chain.crt || { echo "Failed to download CA chain from $CHAIN_URL"; exit 1; }

# Detect OS and install
if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    if [[ "$ID" == "debian" || "$ID" == "ubuntu" || "$ID" == "linuxmint" ]]; then
        sudo cp /tmp/ca-chain.crt /usr/local/share/ca-certificates/certauth.crt
        sudo update-ca-certificates
        echo "✅ CA chain installed (Debian/Ubuntu)"
    elif [[ "$ID" == "centos" || "$ID" == "rhel" || "$ID" == "fedora" ]]; then
        # Install to system trust store (curl, openssl, etc.)
        sudo cp /tmp/ca-chain.crt /etc/pki/ca-trust/source/anchors/certauth.crt
        sudo update-ca-trust

        # Install root CA to NSS database (Firefox, Thunderbird, etc.)
        # The root CA is the self-signed cert (second cert in chain)
        if command -v certutil &>/dev/null; then
            sudo certutil -D -n "CertAuth Root CA" -d sql:/etc/pki/nssdb/ 2>/dev/null || true
            python3 -c "
import re, subprocess
with open('/tmp/ca-chain.crt') as f:
    content = f.read()
certs = re.findall(r'-----BEGIN CERTIFICATE-----.*?-----END CERTIFICATE-----', content, re.DOTALL)
# Find the self-signed root cert (subject == issuer)
for cert in certs:
    subj = subprocess.run(['openssl', 'x509', '-noout', '-subject'], input=cert, capture_output=True, text=True).stdout
    iss = subprocess.run(['openssl', 'x509', '-noout', '-issuer'], input=cert, capture_output=True, text=True).stdout
    if subj.replace('subject=', '') == iss.replace('issuer=', ''):
        with open('/tmp/certauth-root.crt', 'w') as rf:
            rf.write(cert + '\n')
        break
" 2>/dev/null
            if [[ -f /tmp/certauth-root.crt ]]; then
                sudo certutil -A -n "CertAuth Root CA" -t "CT,Cu,Tu" -d sql:/etc/pki/nssdb/ -i /tmp/certauth-root.crt 2>/dev/null
                echo "✅ Root CA imported to NSS database"
                rm -f /tmp/certauth-root.crt
            fi
        fi
        echo "✅ CA chain installed (RHEL/CentOS/Fedora)"
    elif [[ "$ID" == "arch" ]]; then
        sudo cp /tmp/ca-chain.crt /etc/ca-certificates/trust-source/anchors/certauth.crt
        sudo update-ca-trust
        echo "✅ CA chain installed (Arch)"
    elif [[ "$ID" == "alpine" ]]; then
        sudo cp /tmp/ca-chain.crt /usr/local/share/ca-certificates/certauth.crt
        sudo update-ca-certificates
        echo "✅ CA chain installed (Alpine)"
    else
        echo "❌ Unsupported Linux distribution: $ID"
        echo "   Download /tmp/ca-chain.crt and install manually"
        exit 1
    fi
elif [[ "$(uname)" == "Darwin" ]]; then
    sudo security add-trusted-cert -d -r trustRoot -k /Library/Keychains/System.keychain /tmp/ca-chain.crt
    echo "✅ CA chain installed (macOS)"
else
    echo "❌ Unsupported OS: $(uname -s)"
    echo "   Download /tmp/ca-chain.crt and install manually"
    exit 1
fi

# Verify
VERIFY_URL="${CHAIN_URL%/api/ca-chain}"
if curl -sL "$VERIFY_URL" &>/dev/null; then
    echo "🌐 Server at $VERIFY_URL is reachable"
else
    echo "⚠️  Server at $VERIFY_URL is not reachable (expected if not on same network)"
fi

rm -f /tmp/ca-chain.crt
echo "Done!"
'''
    return PlainTextResponse(script, media_type="text/x-shellscript")

@app.get("/setup.ps1")
async def setup_ps1():
    """PowerShell setup script for Windows."""
    script = r'''# CertAuth CA Chain Installer for Windows
# Usage: iex (New-Object Net.WebClient).DownloadString("http://<certauth-ip>/setup.ps1")
#        iwr http://<certauth-ip>/setup.ps1 -UseBasicParsing | iex

param([string]$CertAuthIP = "")

if (-not $CertAuthIP) {
    # Try to detect from environment or prompt
    $CertAuthIP = Read-Host "Enter CertAuth server IP (e.g., 192.168.8.248)"
}

$ChainUrl = "http://$CertAuthIP/api/ca-chain"
$ChainPath = "$env:TEMP\ca-chain.crt"

Write-Host "Downloading CA chain from $ChainUrl ..." -ForegroundColor Cyan
try {
    (New-Object Net.WebClient).DownloadFile($ChainUrl, $ChainPath)
} catch {
    Write-Host "Failed to download CA chain: $_" -ForegroundColor Red
    exit 1
}

# Install to Local Machine Trusted Root store
Write-Host "Installing to Trusted Root Certification Authorities..." -ForegroundColor Cyan
try {
    $store = New-Object System.Security.Cryptography.X509Certificates.X509Store(
        [System.Security.Cryptography.X509Certificates.StoreName]::Root,
        [System.Security.Cryptography.X509Certificates.StoreLocation]::LocalMachine)
    $store.Open([System.Security.Cryptography.X509Certificates.OpenFlags]::ReadWrite)
    $cert = New-Object System.Security.Cryptography.X509Certificates.X509Certificate2($ChainPath)
    $store.Add($cert)
    $store.Close()
    Write-Host "CA chain installed successfully!" -ForegroundColor Green
} catch {
    Write-Host "Failed to install: $_" -ForegroundColor Red
    Write-Host "Run as Administrator and try again." -ForegroundColor Yellow
    exit 1
}

# Verify
try {
    $response = Invoke-WebRequest -Uri "http://$CertAuthIP/api/health" -TimeoutSec 3 -ErrorAction Stop
    Write-Host "Server at http://$CertAuthIP is reachable." -ForegroundColor Green
} catch {
    Write-Host "Server at http://$CertAuthIP is not reachable (expected if not on same network)." -ForegroundColor Yellow
}

Remove-Item $ChainPath -Force -ErrorAction SilentlyContinue
Write-Host "Done!" -ForegroundColor Green
'''
    return PlainTextResponse(script, media_type="text/plain")