from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import FileResponse, JSONResponse
import shutil
import os
import sys
from io import StringIO
import contextlib

# ඔයාගේ original ෆයිල් දෙක මෙතනට import කරනවා
import math_solver
import scan  # අපි ඔයාගේ app.py එක scan.py කියලා rename කළානේ

app = FastAPI()

@app.get("/")
def home():
    return {"message": "Cipher-MD Python Engine is Running! 🚀"}

# =================================================================
# 1️⃣ MATH SOLVER ENDPOINT
# =================================================================
@app.post("/math")
async def solve_math(image: UploadFile = File(...), prompt: str = Form("")):
    temp_filename = f"temp_math_{image.filename}"
    
    # 1. Image එක Save කරගන්නවා
    with open(temp_filename, "wb") as buffer:
        shutil.copyfileobj(image.file, buffer)

    # 2. ඔයාගේ math_solver එකේ Output එක (Print වෙන ඒවා) අල්ලාගන්නවා
    output_capture = StringIO()
    try:
        # Stdout එක capture කරනවා (Console එකට print වෙන දේ variable එකකට ගන්නවා)
        with contextlib.redirect_stdout(output_capture):
            # ඔයාගේ original function එක call කරනවා
            math_solver.solve_math_problem(temp_filename, prompt)
        
        result = output_capture.getvalue()
    except Exception as e:
        result = f"Error: {str(e)}"
    finally:
        # කුණු ෆයිල් අස් කරනවා
        if os.path.exists(temp_filename):
            os.remove(temp_filename)

    # 3. JSON විදිහට උත්තරේ යවනවා
    return {"result": result}

# =================================================================
# 2️⃣ SCAN / REVEAL ENDPOINT
# =================================================================
@app.post("/scan")
async def scan_watermark(image: UploadFile = File(...)):
    in_file = f"in_{image.filename}"
    out_file = f"out_{image.filename}"

    # 1. Image එක Save කරගන්නවා
    with open(in_file, "wb") as buffer:
        shutil.copyfileobj(image.file, buffer)

    try:
        # 2. ඔයාගේ scan script (reveal_watermark) එක run කරනවා
        scan.reveal_watermark(in_file, out_file)

        # 3. හැදුනු ෆොටෝ එක ආපහු යවනවා
        if os.path.exists(out_file):
            return FileResponse(out_file)
        else:
            return {"error": "Processing failed"}
            
    except Exception as e:
        return {"error": str(e)}
        
    finally:
        # Input ෆයිල් එක විතරක් මකනවා (Output එක යැව්වට පස්සේ මැකෙන විදිහට HuggingFace එකේ auto clean වෙනවා)
        if os.path.exists(in_file):
            os.remove(in_file)

# Docker එකේ Run වෙන්න ඕන නිසා
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)