from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from io import BytesIO
import asyncio
import tempfile
import zipfile
import os

app = FastAPI()

@app.post("/tex2pdf")
async def convert_to_pdf(zip_file: UploadFile = File(...)):
    if zip_file.filename.endswith('.zip'):
        with tempfile.TemporaryDirectory() as tmpdirname:
            # Unpack the zip file
            with zipfile.ZipFile(BytesIO(await zip_file.read())) as z:
                z.extractall(tmpdirname)
        
            # Change working directory to tmpdirname
            os.chdir(tmpdirname)
            
            # Find the main LaTeX file (assuming a convention, e.g., main.tex)
            main_tex_file = 'main.tex'
            main_tex_path = os.path.join(tmpdirname, main_tex_file)
            if not os.path.exists(main_tex_path):
                raise HTTPException(status_code=400, detail="Main LaTeX file (main.tex) not found in the zip.")

            # Compile the LaTeX document
            cmd = ['pdflatex', '-interaction=nonstopmode', '-output-directory', tmpdirname, main_tex_path]
            process = await asyncio.create_subprocess_exec(*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
            try:
                print(f"Running pdflatex on {main_tex_path}")
                stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=120)
                # print(f"pdflatex output: {stdout.decode()}")
                # print(f"pdflatex errors: {stderr.decode()}")
            except asyncio.TimeoutError:
                return {"error": "Conversion timed out."}


            if process.returncode != 0:
                # Compilation failed
                return {
                    "error": "Conversion failed.",
                    "details": {
                        "stderr": stderr.decode(),
                        "stdout": stdout.decode(),
                    },
                }

            # Assuming the output PDF has the same base name as the main LaTeX file
            output_pdf_path = os.path.join(tmpdirname, 'main.pdf')
            if os.path.exists(output_pdf_path):
                with open(output_pdf_path, 'rb') as f:
                    pdf_content = f.read()
                return StreamingResponse(BytesIO(pdf_content), media_type='application/pdf')
            else:
                return {"error": "PDF file not generated."}
    else:
        raise HTTPException(status_code=400, detail="Uploaded file is not a zip file.")