app.py

import os
import uuid
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException, File, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from playwright.async_api import async_playwright
from huggingface_hub import HfApi, upload_file


class HTMLRequest(BaseModel):
    html_content: str


class PDFResponse(BaseModel):
    pdf_url: str
    filename: str
    message: str
    repository_url: str


async def html_to_pdf_api(html_content: str) -> tuple[str, str]:
    """
    HTMLコンテンツをPDFに変換してHugging Faceデータセットリポジトリにアップロード
    knowledge.txtのPlaywright手法を使用（Async版）
    """
    try:
        # 環境変数からリポジトリ情報を取得
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        hf_token = os.getenv("HF_TOKEN")
        
        if not hf_repo_id:
            raise HTTPException(status_code=500, detail="HF_DATASET_REPO_ID環境変数が設定されていません")
        if not hf_token:
            raise HTTPException(status_code=500, detail="HF_TOKEN環境変数が設定されていません")
        
        # 一意のファイル名を生成
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"document_{timestamp}_{unique_id}.pdf"
        
        # 一時ファイルでPDFを生成
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
            temp_path = temp_file.name
            
            # Playwrightでヘッドレスブラウザを起動（Async版）
            async with async_playwright() as pw:
                browser = await pw.chromium.launch(headless=True)
                page = await browser.new_page()
                
                # HTMLコンテンツを設定（外部リソース読み込み待機）
                await page.set_content(html_content, wait_until="networkidle")
                
                # 印刷メディアを有効にする
                await page.emulate_media(media="print")
                
                # PDFを生成（test.htmlの設定に準拠）
                await page.pdf(
                    path=temp_path,
                    format="A4",
                    print_background=True,
                    margin={"top":"15mm","bottom":"15mm","left":"15mm","right":"15mm"},
                    scale=0.9  # 90%に縮小してA4 2ページに収める
                )
                
                await browser.close()
            
            # Hugging Face リポジトリにアップロード
            api = HfApi(token=hf_token)
            
            upload_file(
                path_or_fileobj=temp_path,
                path_in_repo=f"pdfs/{filename}",
                repo_id=hf_repo_id,
                repo_type="dataset",
                token=hf_token,
                commit_message=f"Add PDF: {filename}"
            )
            
            # 一時ファイルを削除
            os.unlink(temp_path)
            
            # ダウンロードURLを生成
            pdf_url = f"https://huggingface.co/datasets/{hf_repo_id}/resolve/main/pdfs/{filename}"
            
            return filename, pdf_url
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"PDF生成中にエラーが発生しました: {str(e)}")


# FastAPIアプリケーションの初期化
app = FastAPI(
    title="HTML to PDF Converter API",
    description="日本語対応のHTML→PDF変換API。複雑なレイアウトやWebフォントを含むHTMLコンテンツを、正確にA4サイズのPDFに変換します。",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# CORS設定
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def root():
    """
    API情報を返すルートエンドポイント
    """
    hf_repo_id = os.getenv("HF_DATASET_REPO_ID", "未設定")
    return {
        "message": "HTML to PDF Converter API",
        "description": "HTMLコンテンツをA4サイズのPDFに変換し、Hugging Faceデータセットリポジトリに保存するAPI",
        "version": "1.0.0",
        "storage": f"Hugging Face Dataset Repository: {hf_repo_id}",
        "endpoints": {
            "convert": "/convert - HTMLをPDFに変換してHFリポジトリに保存",
            "files": "/files - HFリポジトリ内のPDFファイル一覧",
            "docs": "/docs - API仕様書",
            "health": "/health - ヘルスチェック"
        }
    }


@app.post("/convert", response_model=PDFResponse)
async def convert_html_to_pdf(request: HTMLRequest):
    """
    HTMLコンテンツをPDFに変換してHugging Faceデータセットリポジトリに保存するエンドポイント
    
    - **html_content**: 変換するHTMLコンテンツ
    
    Returns:
    - **pdf_url**: 生成されたPDFのダウンロードURL（Hugging Face上）
    - **filename**: PDFファイル名
    - **message**: 処理結果メッセージ
    - **repository_url**: リポジトリURL
    """
    if not request.html_content or not request.html_content.strip():
        raise HTTPException(status_code=400, detail="HTMLコンテンツが空です")
    
    try:
        filename, pdf_url = await html_to_pdf_api(request.html_content)
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        repository_url = f"https://huggingface.co/datasets/{hf_repo_id}"
        
        return PDFResponse(
            pdf_url=pdf_url,
            filename=filename,
            message=f"✅ PDFが正常に生成され、Hugging Faceリポジトリに保存されました: {filename}",
            repository_url=repository_url
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"処理中にエラーが発生しました: {str(e)}")


@app.get("/download/{filename}")
async def download_pdf(filename: str):
    """
    Hugging FaceリポジトリのPDFファイルへリダイレクトするエンドポイント
    
    - **filename**: ダウンロードするPDFファイル名
    """
    hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
    if not hf_repo_id:
        raise HTTPException(status_code=500, detail="HF_DATASET_REPO_ID環境変数が設定されていません")
    
    # Hugging Face上のファイルURLにリダイレクト
    pdf_url = f"https://huggingface.co/datasets/{hf_repo_id}/resolve/main/pdfs/{filename}"
    return RedirectResponse(url=pdf_url)


@app.get("/health")
async def health_check():
    """
    ヘルスチェックエンドポイント
    """
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}


@app.get("/files")
async def list_files():
    """
    Hugging Faceリポジトリ内のPDFファイル一覧を取得するエンドポイント
    """
    try:
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        hf_token = os.getenv("HF_TOKEN")
        
        if not hf_repo_id:
            raise HTTPException(status_code=500, detail="HF_DATASET_REPO_ID環境変数が設定されていません")
        if not hf_token:
            raise HTTPException(status_code=500, detail="HF_TOKEN環境変数が設定されていません")
        
        api = HfApi(token=hf_token)
        
        # リポジトリ内のファイル一覧を取得
        try:
            repo_files = api.list_repo_files(repo_id=hf_repo_id, repo_type="dataset")
            pdf_files = [f for f in repo_files if f.startswith("pdfs/") and f.endswith(".pdf")]
            
            files = []
            for file_path in pdf_files:
                filename = Path(file_path).name
                file_info = api.get_paths_info(repo_id=hf_repo_id, paths=[file_path], repo_type="dataset")[0]
                
                files.append({
                    "filename": filename,
                    "path": file_path,
                    "size": file_info.size if hasattr(file_info, 'size') else 0,
                    "last_modified": file_info.last_commit.date.isoformat() if hasattr(file_info, 'last_commit') and file_info.last_commit else None,
                    "download_url": f"https://huggingface.co/datasets/{hf_repo_id}/resolve/main/{file_path}",
                    "api_download_url": f"/download/{filename}"
                })
            
            return {
                "repository": hf_repo_id,
                "total_files": len(files),
                "files": files
            }
            
        except Exception as e:
            # リポジトリが空の場合やpdfsフォルダが存在しない場合
            return {
                "repository": hf_repo_id,
                "total_files": 0,
                "files": [],
                "note": "No PDF files found or repository is empty"
            }
            
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"ファイル一覧取得中にエラーが発生しました: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)