from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from datetime import datetime
import json
from huggingface_hub import HfApi
import os
from huggingface_hub import hf_hub_download

api = HfApi()
app = FastAPI() 

modelName = "MitchellKil/gaelic-ipa-byt5" # my fine tuned model
feedbackFile = "mitchellkil/gaelic-app-feedback" # file to write feedback to

HF_TOKEN = os.environ["HF_TOKEN"] # set in secrets


print("Loading IPA model")

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(modelName)
model = AutoModelForSeq2SeqLM.from_pretrained(modelName).to(device)
model.eval()

print("Model loaded") 

class gaelicText(BaseModel):
    text: str


def gaelicToIPA(text: str) -> str:
    # Input text - > target text for trained model
    prompt = f"convert Gaelic to IPA: {text}"

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=110
    ).to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=False
        )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return decoded.strip()

@app.post("/predict") # call to  return text to project
def predict(request: gaelicText):
    ipa = gaelicToIPA(request.text)
    return {"ipa": ipa}

class Feedback(BaseModel):
    type: str
    feedback: str
    input_text: str
    translation: str
    phonetic: str


@app.post("/feedback")
def save_feedback(data: Feedback):
    feedbackText = {
        "type": data.type,
        "timestamp": datetime.utcnow().isoformat(),
        "input_text": data.input_text,
        "translation": data.translation,
        "phonetic": data.phonetic,
        "feedback": data.feedback
    }
    try: # current feedback file
        local_file = hf_hub_download(
            repo_id=dataset,
            filename="feedback.jsonl",
            repo_type="dataset",
            token=HF_TOKEN
        )
        with open(local_file, "r", encoding="utf-8") as f:
            lines = f.readlines()
    except Exception:
        lines = []

    lines.append(json.dumps(feedbackText) + "\n") # append

    tempPath = "feedbackTemp.jsonl" # create temp file
    with open(temp_path, "w", encoding="utf-8") as f:
        f.writelines(lines)

    api.upload_file( # reupload full file
        path_or_fileobj=tempPath,
        path_in_repo="feedback.jsonl",
        repo_id=dataset,
        repo_type="feedbackFile",
        token=HF_TOKEN,
        commit_message="New feedback added"
    )

    return {"status": "Thanks for your feedback"}