"""Qwen-powered speech coaching service."""

import modal

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
CACHE_DIR = "/cache"
UNMEASURED_TERMS = (
    "pause",
    "breath",
    "tone",
    "emotion",
    "pronunciation",
    "emphasis",
    "eye contact",
    "volume",
)

app = modal.App("gtrox-coach")
model_cache = modal.Volume.from_name("gtrox-model-cache", create_if_missing=True)

image = (
    modal.Image.debian_slim(python_version="3.12")
    .env({"HF_HOME": CACHE_DIR, "HF_XET_HIGH_PERFORMANCE": "1"})
    .uv_pip_install(
        "torch==2.7.1",
        "transformers==4.53.3",
        "accelerate==1.8.1",
        "huggingface_hub[hf-xet]==0.33.4",
        "fastapi[standard]==0.115.4",
    )
)


@app.cls(
    gpu="T4",
    image=image,
    timeout=300,
    max_containers=1,
    scaledown_window=300,
    volumes={CACHE_DIR: model_cache},
)
class SpeechCoach:
    @modal.enter()
    def setup(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float16,
            device_map="auto",
        )

    @modal.fastapi_endpoint(method="POST")
    def coach(self, payload: dict):
        import torch

        transcript = str(payload.get("transcript", ""))[:8000]
        metrics = payload.get("metrics", {})
        messages = [
            {
                "role": "system",
                "content": (
                    "You are a supportive professional speech coach. Give concise, "
                    "specific advice grounded only in the provided transcript and "
                    "metrics. Do not invent pause, tone, emotion, or pronunciation "
                    "observations. Interpret speaking pace as: below 110 WPM is slow, "
                    "110-180 WPM is conversational, and above 180 WPM is fast. Treat "
                    "higher filler and repetition density as improvement areas. Never "
                    "contradict the supplied numbers. You may discuss only measured "
                    "pace, fillers, repetitions, and the transcript's wording or "
                    "organization. Never mention pauses, breathing, tone, emotion, "
                    "pronunciation, emphasis, eye contact, or volume, even as advice. "
                    "Return markdown with headings: Summary, Strengths, Improve Next, "
                    "and Practice Exercise. Give no more than three recommendations."
                ),
            },
            {
                "role": "user",
                "content": f"Metrics: {metrics}\n\nTranscript:\n{transcript}",
            },
        ]
        prompt = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda")

        with torch.inference_mode():
            generated = self.model.generate(
                **inputs,
                max_new_tokens=240,
                do_sample=False,
                repetition_penalty=1.05,
            )

        new_tokens = generated[0][inputs["input_ids"].shape[1] :]
        coaching = self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
        coaching = "\n".join(
            line
            for line in coaching.splitlines()
            if not any(term in line.lower() for term in UNMEASURED_TERMS)
        ).strip()

        if len(coaching) < 120:
            coaching = (
                "## Summary\n"
                f"Your measured speaking pace is {metrics.get('wpm', 0)} WPM with a "
                f"fluency score of {metrics.get('fluency_score', 0)}/100.\n\n"
                "## Strengths\n"
                f"- Filler count: {metrics.get('filler_count', 0)}\n"
                f"- Immediate repetitions: {metrics.get('repetitions', 0)}\n\n"
                "## Improve Next\n"
                "- Organize the message around one clear main idea and supporting points.\n\n"
                "## Practice Exercise\n"
                "- Deliver the same message with an opening, three points, and a conclusion."
            )

        return {"coaching": coaching, "model": MODEL_NAME}