from datetime import datetime, timezone

from fastapi import APIRouter, Request
from groq import Groq, RateLimitError
from pydantic import BaseModel

from src.api.utils import (
    check_user_rate_limit,
    get_groq_keys,
    track_key_usage,
    user_request_tracker,
    verify_user_license,
)
from src.common.logger import logger
from src.common.utils import response_error, response_success

router = APIRouter()


class TranslationRequest(BaseModel):
    email: str
    license_key: str
    texts: list[str]
    source_language: str = "auto"
    target_language: str = "en"
    custom_prompt: str = None


@router.post("/translate")
async def translate_texts(req: TranslationRequest, request: Request):
    MODEL = "openai/gpt-oss-120b"
    _, error = verify_user_license(req.email, req.license_key)
    if error:
        return response_error(error, f"License verification failed: {error}", 403)

    if not check_user_rate_limit(req.email):
        remaining = user_request_tracker[req.email]["reset_at"] - datetime.now(
            timezone.utc
        )
        return response_error(
            "USER_RATE_LIMIT",
            f"Rate limit exceeded. Try again in {int(remaining.total_seconds())} seconds",
            429,
        )

    if not req.texts:
        return response_error("INVALID_INPUT", "texts array cannot be empty", 400)

    api_keys = get_groq_keys()
    if not api_keys:
        return response_error("NO_API_KEYS", "No Groq API keys configured", 503)

    logger.info(f"Translation request received: {req.email}, {len(req.texts)} texts")
    base_instruction = (
        "You are a professional translator. First, correct any typos or grammatical errors in the input, "
        "then translate it naturally, accurately, and contextually. Do not explain unless asked."
    )
    if req.custom_prompt:
        base_instruction += f" {req.custom_prompt}"

    numbered_texts = "\n".join(f"{i+1}. {t}" for i, t in enumerate(req.texts))

    if req.source_language == "auto":
        prompt = (
            f"{base_instruction}\n\n"
            f"Translate the following texts to {req.target_language}.\n"
            f"Return exactly {len(req.texts)} lines, numbered.\n\n"
            f"{numbered_texts}"
        )
    else:
        prompt = (
            f"{base_instruction}\n\n"
            f"Translate from {req.source_language} to {req.target_language}.\n"
            f"Return exactly {len(req.texts)} lines, numbered.\n\n"
            f"{numbered_texts}"
        )

    last_error = None
    max_retries = 3

    for retry in range(max_retries):
        for i, api_key in enumerate(api_keys):
            try:
                client = Groq(api_key=api_key)

                completion = client.chat.completions.create(
                    model=MODEL,
                    messages=[
                        {"role": "system", "content": base_instruction},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=1,
                    top_p=1,
                    max_completion_tokens=8192,
                )
                track_key_usage(api_key)

                output = completion.choices[0].message.content.strip()
                lines = [
                    l.split(". ", 1)[1] if ". " in l else l
                    for l in output.split("\n")
                    if l.strip()
                ]

                if len(lines) != len(req.texts):
                    if retry < max_retries - 1:
                        prompt += (
                            f"\n\nIMPORTANT: Return EXACTLY {len(req.texts)} items."
                        )
                        continue

                return response_success(
                    {
                        "original_texts": req.texts,
                        "translated_texts": lines,
                        "source_language": req.source_language,
                        "target_language": req.target_language,
                        "model": MODEL,
                    }
                )

            except RateLimitError:
                continue

            except Exception as e:
                last_error = str(e)
                continue

    return response_error("TRANSLATION_FAILED", f"All keys failed: {last_error}", 500)