from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer
import traceback
import whisper
import librosa
import numpy as np
import torch
import uvicorn
import base64
import io
from voxcpm import VoxCPM

asr_model = whisper.load_model("models/wpt/wpt.pt")
model_name = "models/Llama-3.2-1B-Instruct"
tok = AutoTokenizer.from_pretrained(model_name)
lm = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
).eval()

tts = VoxCPM.from_pretrained(
    "models/VoxCPM-0.5B", 
    local_files_only=True, 
    load_denoiser=True, 
    zipenhancer_model_id="models/iic/speech_zipenhancer_ans_multiloss_16k_base"
    )

def chat(system_prompt: str, user_prompt: str) -> str:
    """
    Run one turn of chat with a system + user message.
    Extra **gen_kwargs are forwarded to `generate()`.
    """
    print("LLM init...")
    messages = [
        {"role": "system",    "content": system_prompt},
        {"role": "user",      "content": user_prompt},
    ]

    # `add_generation_prompt=True` automatically appends the
    #   <|start_header_id|>assistant … header so the model knows to respond.
    # Get both input_ids and attention_mask
    inputs = tok.apply_chat_template(
        messages, 
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True  # Returns dict with input_ids and attention_mask
    )
    
    # Move to device
    input_ids = inputs["input_ids"].to(lm.device)
    attention_mask = inputs["attention_mask"].to(lm.device)

    with torch.inference_mode():
        output_ids = lm.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,  # Proper attention mask
            pad_token_id=tok.eos_token_id,  # Explicit pad token
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.2,
            repetition_penalty=1.1,
            top_k=100,
            top_p=0.95,
        )

    # Strip the prompt part and return only the newly-generated answer
    answer = tok.decode(
        output_ids[0][input_ids.shape[-1]:],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True,
    )
    print("LLM answer done.")
    return answer.strip()

def gt(audio: np.ndarray, sr: int):
    print("Starting ASR transcription...")
    ss = audio.squeeze().astype(np.float32)
    if sr != 16_000:
        ss = librosa.resample(audio, orig_sr=sr, target_sr=16_000)

    result = asr_model.transcribe(ss, fp16=False, language=None)
    transcribed_text = result["text"].strip()
    print(f"ASR done. Transcribed: '{transcribed_text}'")
    return transcribed_text


def sample(rr: str) -> str:
    if rr.strip() == "":
        rr = "Hello "

    inputs = tok(rr, return_tensors="pt").to(lm.device)

    with torch.inference_mode():
        out_ids = lm.generate(
            **inputs,
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.2,
            repetition_penalty=1.1,
            top_k=100,
            top_p=0.95,
        )

    return tok.decode(
        out_ids[0][inputs.input_ids.shape[-1] :], skip_special_tokens=True
    )


INITIALIZATION_STATUS = {"model_loaded": True, "error": None}


class GenerateRequest(BaseModel):
    audio_data: str = Field(
        ...,
        description="",
    )
    sample_rate: int = Field(..., description="")


class GenerateResponse(BaseModel):
    audio_data: str = Field(..., description="")


app = FastAPI(title="V1", version="0.1")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def b64(b64: str) -> np.ndarray:
    raw = base64.b64decode(b64)
    return np.load(io.BytesIO(raw), allow_pickle=False)


def ab64(arr: np.ndarray, sr: int) -> str:
    buf = io.BytesIO()
    resampled = librosa.resample(arr, orig_sr=16000, target_sr=sr)
    np.save(buf, resampled.astype(np.float32))
    return base64.b64encode(buf.getvalue()).decode()

@app.get("/api/v1/health")
def health_check():
    """Health check endpoint"""
    status = {
        "status": "healthy",
        "model_loaded": INITIALIZATION_STATUS["model_loaded"],
        "error": INITIALIZATION_STATUS["error"],
    }
    return status


@app.post("/api/v1/v2v", response_model=GenerateResponse)
def generate_audio(req: GenerateRequest):
    print("=== V2V Request Started ===")
    audio_np = b64(req.audio_data)
    if audio_np.ndim == 1:
        audio_np = audio_np.reshape(1, -1)
    print(f"Audio shape: {audio_np.shape}, Sample rate: {req.sample_rate}")

    system_prompt = "You are a helpful assistant who tries to help answer the user's question. This is a part of voice assistant system, don't generate anything other than pure text, no symbols other than simple punctuation marks."

    try:
        # Step 1: ASR
        text = gt(audio_np, req.sample_rate)

        # Step 2: LLM
        response_text = chat(system_prompt, user_prompt=text)
        print(f"LLM response len chars: '{len(response_text)}'")
        print(f"LLM response: '{response_text}'")

        import time
        start_time = time.perf_counter()
        audio_out = tts.generate(
            text=response_text,
            prompt_wav_path=None,      # optional: path to a prompt speech for voice cloning
            prompt_text=None,          # optional: reference text
            cfg_value=2.0,             # LM guidance on LocDiT, higher for better adherence to the prompt, but maybe worse
            inference_timesteps=10,   # LocDiT inference timesteps, higher for better result, lower for fast speed
            normalize=True,           # enable external TN tool
            denoise=True,             # enable external Denoise tool
            retry_badcase=True,        # enable retrying mode for some bad cases (unstoppable)
            retry_badcase_max_times=3,  # maximum retrying times
            retry_badcase_ratio_threshold=6.0, # maximum length restriction for bad case detection (simple but effective), it could be adjusted for slow pace speech
        )
        print("TTS generation complete.")
        end_time = time.perf_counter()
        print(f"TTS generation took {end_time - start_time:.2f} seconds.")
        print("=== V2V Request Complete ===")
    except Exception as e:
        print(f"ERROR in V2V: {e}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"{e}")

    return GenerateResponse(audio_data=ab64(audio_out, req.sample_rate))

@app.post("/api/v1/v2t")
def generate_text(req: GenerateRequest):
    audio_np = b64(req.audio_data)
    if audio_np.ndim == 1:
        audio_np = audio_np.reshape(1, -1)

    try:
        text = gt(audio_np, req.sample_rate)
        print(f"Transcribed text: {text}")
        # response_text = sample(text)
        system_prompt = "You are a helpful assistant who tries to help answer the user's question."
        response_text = chat(system_prompt, user_prompt=text)
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"{e}")

    return {"text": response_text}


if __name__ == "__main__":
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=False)