import os
from fastapi import FastAPI, UploadFile, File
import fitz
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import json

app = FastAPI()

# --- NEW: Download logic to bypass 1GB repo limit ---
REPO_ID = "amielitos/text-To-JSON" # Change this to your Model Repo ID
FILENAME = "phi-3.5-mini-instruct.Q4_K_M.gguf"

# This downloads the file to a local cache folder and returns the path
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

# Now load the model from that path
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=4
    n_gpu_layers=0
)

def extract_context(pdf_bytes):
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    text = ""
    # Extract from first 3 pages for context
    for i in range(min(3, len(doc))):
        text += doc[i].get_text()
    return text[:2000]

@app.post("/translate")
async def translate_pdf(file: UploadFile = File(...)):
    pdf_content = await file.read()
    context = extract_context(pdf_content)
    
    prompt = f"<|user|>\nSummarize the following scientific text and output a multiple-choice question in JSON format.\n{context}<|end|>\n<|assistant|>\n"
    
    output = llm(prompt, max_tokens=512, stop=["<|end|>"], temperature=0.1)
    
    response_text = output["choices"][0]["text"].strip()
    try:
        start = response_text.find("{")
        end = response_text.rfind("}") + 1
        return json.loads(response_text[start:end])
    except:
        return {"error": "JSON parse error", "raw_text": response_text}

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