import os
import re  # Import the regular expression module
from typing import List

from fastapi import FastAPI, HTTPException, status
from pydantic import BaseModel
import requests
from fastapi.responses import HTMLResponse

app = FastAPI()

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise ValueError("The GEMINI_API_KEY environment variable is not set.")

GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-lite:generateContent?key=" + GEMINI_API_KEY


class TranslationRequest(BaseModel):
    text: str
    target_language: str
    source_language: str = None


class TranslationResponse(BaseModel):  # For /detect_language
    translated_text: str
    source_language: str
    target_language: str

# --- Helper Functions ---

def detect_language_and_options(text: str):
    """Detects the language and provides translation options."""
    prompt = f"""Please identify the language of the text provided and then offer translation options as numbered choices (1-5).  Use this format:  "The text is in [Language].  Choose a language to translate to: 1. [Option 1], 2. [Option 2], 3. [Option 3], 4. [Option 4], 5. [Option 5]"

    Input Text: {text}"""

    request_data = {
        "contents": [{
            "role": "user",
            "parts": [{"text": prompt}]
        }]
    }

    try:
        response = requests.post(GEMINI_API_URL, json=request_data)
        response.raise_for_status()
        response_json = response.json()

        try:
            response_text = response_json['candidates'][0]['content']['parts'][0]['text']
            source_language = response_text.split("The text is in ")[1].split(".")[0].strip()
            options_str = response_text.split("Choose a language to translate to:")[1].strip()
            options_list = [opt.split(". ")[1].strip() for opt in options_str.split(", ")]
            while len(options_list) < 5:
                options_list.append("Option Not Available")
            options_list = options_list[:5]
            options = {str(i + 1): lang for i, lang in enumerate(options_list)}
            return source_language, options
        except (KeyError, IndexError, AttributeError) as e:
            raise HTTPException(status_code=500, detail=f"Error parsing Gemini API response: {e}")
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Error communicating with Gemini API: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")


def translate_with_gemini(text: str, source_language: str, target_language: str) -> str:
    """Translates text, requesting HTML, and removes Markdown code blocks."""

    prompt = f"""Translate the following text from {source_language} to {target_language}. Return *only* the translated text.  Format the response as HTML. If any part of the translated text should be emphasized (like an example or key phrase), use `<strong>` tags around that part. Do not include ANY additional text like 'Here is the translation:' or similar.  Only the translated text, formatted as HTML.

    Text:
    {text}"""

    request_data = {
        "contents": [{
            "role": "user",
            "parts": [{"text": prompt}]
        }]
    }

    try:
        response = requests.post(GEMINI_API_URL, json=request_data)
        response.raise_for_status()
        response_json = response.json()

        try:
            translated_text = response_json['candidates'][0]['content']['parts'][0]['text']

            # Remove Markdown code blocks (```html ... ```)
            translated_text = re.sub(r"```html\n?([\s\S]*?)\n?```", r"\1", translated_text).strip()
            translated_text = re.sub(r"```([\s\S]*?)```", r"\1", translated_text).strip() # Remove plain ``` ... ```


            return translated_text
        except (KeyError, IndexError) as e:
            raise HTTPException(status_code=500, detail=f"Error parsing Gemini API response: {e}")
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Error communicating with Gemini API: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")


@app.post("/translate", response_class=HTMLResponse, status_code=status.HTTP_200_OK)
async def translate(request: TranslationRequest):
    """Translates text, returning HTML."""

    if not request.text:
        raise HTTPException(status_code=400, detail="Text to translate cannot be empty.")
    if not request.target_language:
        raise HTTPException(status_code=400, detail="Target language must be provided.")

    if request.source_language:
        source_language = request.source_language
    else:
        try:
            source_language, _ = detect_language_and_options(request.text)
        except HTTPException as e:
            raise e

    supported_languages = ["Afrikaans", "Arabic", "Armenian", "Azerbaijani", "Belarusian", "Bosnian", "Bulgarian",
                           "Catalan", "Chinese", "Croatian", "Czech", "Danish", "Dutch", "English", "Estonian",
                           "Finnish", "French", "Galician", "German", "Greek", "Hebrew", "Hindi", "Hungarian",
                           "Icelandic", "Indonesian", "Italian", "Japanese", "Kannada", "Kazakh", "Korean", "Latvian",
                           "Lithuanian", "Macedonian", "Malay", "Marathi", "Maori", "Nepali", "Norwegian", "Persian",
                           "Polish", "Portuguese", "Romanian", "Russian", "Serbian", "Slovak", "Slovenian", "Spanish",
                           "Swahili", "Swedish", "Tagalog", "Tamil", "Thai", "Turkish", "Ukrainian", "Urdu",
                           "Vietnamese", "Welsh"]

    if request.target_language not in supported_languages:
        raise HTTPException(status_code=400, detail=f"Target language '{request.target_language}' is not supported.")

    try:
        translated_html = translate_with_gemini(request.text, source_language, request.target_language)
        return translated_html
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")



@app.post("/detect_language", status_code=status.HTTP_200_OK)
async def detect_language(text: str = ""):
    """Detects the language of the input text and provides translation options."""
    if not text:
        raise HTTPException(status_code=400, detail="Text to detect cannot be empty.")
    try:
        source_language, options = detect_language_and_options(text)
        return {"source_language": source_language, "translation_options": options}
    except HTTPException as e:
        raise e
    except Exception as e:
         raise HTTPException(status_code=500,detail=f"An unexpected error occurred {e}")