from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from pdf2image import convert_from_path
import os
import shutil
import subprocess
import sys
import json
from pathlib import Path
import traceback

app = FastAPI(title="OMRChecker FastAPI")

# Base directories
HF_HOME = "/tmp/huggingface"
INPUT_DIR = f"{HF_HOME}/inputs"
OUTPUT_DIR = f"{HF_HOME}/outputs"

# Ensure directories exist
os.environ["HF_HOME"] = HF_HOME
os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

@app.post("/process-omr")
async def process_omr(
    config: UploadFile = File(...),
    evaluation: UploadFile = File(...),
    template: UploadFile = File(...),
    omr_marker: UploadFile = File(...),
    pdf_file: UploadFile = File(...),
):


    try:
        print("Cleaning old directories...")
        shutil.rmtree(INPUT_DIR, ignore_errors=True)
        shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
        os.makedirs(INPUT_DIR, exist_ok=True)
        os.makedirs(OUTPUT_DIR, exist_ok=True)

        print("Saving uploaded files to INPUT_DIR...")
        uploaded_files = {
            "config.json": config,
            "evaluation.json": evaluation,
            "template.json": template,
            "omr_marker.jpg": omr_marker,
            pdf_file.filename: pdf_file,
        }

        for filename, file in uploaded_files.items():
            file_path = os.path.join(INPUT_DIR, filename)
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            print(f"✅ Saved: {filename}")

        print("Converting PDF to images...")
        pdf_path = os.path.join(INPUT_DIR, pdf_file.filename)
        images = convert_from_path(pdf_path)
        image_paths = []

        for i, img in enumerate(images):
            img_filename = f"converted_page_{i+1}.jpg"
            img_output_path = os.path.join(INPUT_DIR, img_filename)
            img.save(img_output_path, "JPEG")
            image_paths.append(img_output_path)
            print(f"🖼️ Converted and saved: {img_output_path}")

        script_path = os.path.join(os.path.dirname(__file__), "main.py")

        process = subprocess.run(
            [sys.executable, script_path, "--inputDir", INPUT_DIR, "--outputDir", OUTPUT_DIR],
            capture_output=True,
            text=True,
        )

        print(process.stdout)
        print(process.stderr)

        response_json = os.path.join(OUTPUT_DIR, "read_response.json")
        score_txt = os.path.join(OUTPUT_DIR, "score.txt")

        results = {"converted_images": [], "read_response": None, "score": None}

        output_images = [
            os.path.join(OUTPUT_DIR, f)
            for f in os.listdir(OUTPUT_DIR)
            if f.endswith(".jpg")
        ]

        for img_path in output_images:
            results["converted_images"].append(img_path)

        # Load JSON and score files
        if os.path.exists(response_json):
            with open(response_json, "r") as f:
                results["read_response"] = json.load(f)
            print("Loaded read_response.json")

        if os.path.exists(score_txt):
            with open(score_txt, "r") as f:
                results["score"] = f.read().strip()
            print("Loaded score.txt")

        print("Processing complete.")
        return JSONResponse(content=results)

    except Exception as e:
        print("Exception during OMR processing:")
        traceback.print_exc()
        return JSONResponse(content={"error": str(e)}, status_code=500)