import os
from fastapi import FastAPI, Request, Form, File, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from groq import Groq
import io

# Set up the Groq client
os.environ["GROQ_API_KEY"] = "your_groq_api_key_here"
client = Groq(api_key=os.environ["GROQ_API_KEY"])

# Initialize FastAPI app and template engine
app = FastAPI()
templates = Jinja2Templates(directory="templates")

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/transcribe")
async def transcribe_audio(audio_data: UploadFile = File(...), language: str = Form(...)):
    try:
        audio_content = await audio_data.read()
        # Transcribe the audio based on the selected language
        transcription = client.audio.transcriptions.create(
            file=(audio_data.filename, audio_content),
            model="whisper-large-v3",
            prompt="Transcribe the audio accurately based on the selected language.",
            response_format="text",
            language=language,
        )

        return JSONResponse(content={'transcription': transcription})

    except Exception as e:
        return JSONResponse(status_code=500, content={'error': str(e)})

@app.post("/check_grammar")
async def check_grammar(transcription: str = Form(...), language: str = Form(...)):
    if not transcription or not language:
        return JSONResponse(status_code=400, content={'error': 'Missing transcription or language selection'})

    try:
        # Grammar check
        grammar_prompt = (
            f"Briefly check the grammar of the following text in {language}: {transcription}. "
            "Identify any word that does not belong to the selected language and flag it. "
            "Based on the number of incorrect words also check the grammar deeply and carefully. "
            "Provide a score from 1 to 10 based on the grammar accuracy, reducing points for incorrect words and make sure to output the score on a new line after two line breaks like 'SCORE='."
        )

        grammar_check_response = client.chat.completions.create(
            model="llama3-groq-70b-8192-tool-use-preview",
            messages=[{"role": "user", "content": grammar_prompt}]
        )
        grammar_feedback = grammar_check_response.choices[0].message.content.strip()

        # Vocabulary check
        vocabulary_prompt = (
            f"Check the vocabulary accuracy of the following text in {language}: {transcription}. "
            "Identify any word that does not belong to the selected language and flag it. "
            "Based on the number of incorrect words also check the grammar deeply and carefully. "
            "Provide a score from 1 to 10 based on the vocabulary accuracy reducing points for incorrect words and make sure to output the score on a new line after two line breaks like 'SCORE='."
        )

        vocabulary_check_response = client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=[{"role": "user", "content": vocabulary_prompt}]
        )
        vocabulary_feedback = vocabulary_check_response.choices[0].message.content.strip()

        # Calculate scores
        grammar_score = calculate_score(grammar_feedback)
        vocabulary_score = calculate_score(vocabulary_feedback)

        return JSONResponse(content={
            'grammar_feedback': grammar_feedback,
            'vocabulary_feedback': vocabulary_feedback,
            'grammar_score': grammar_score,
            'vocabulary_score': vocabulary_score
        })

    except Exception as e:
        return JSONResponse(status_code=500, content={'error': str(e)})


def calculate_score(feedback: str) -> int:
    """ 
    Calculate score based on feedback content.
    This function searches for the keyword 'SCORE=' or similar variations 
    (SCORE:, score:, etc.) and extracts the score value.
    """
    import re
    match = re.search(r'(SCORE=|score=|SCORE:|score:|SCORE = )\s*(\d+)', feedback)
    
    if match:
        return int(match.group(2))
    
    return 0  # Return default score of 0 if no score is found


# Run the FastAPI app (only needed for local development)
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)