import io, json, logging, os, time, uuid, asyncio
from collections import defaultdict, deque

import numpy as np
import soundfile as sf
import librosa
import ctranslate2
import torch
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Depends, Request, WebSocket, WebSocketDisconnect, Query
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import JSONResponse
from huggingface_hub import HfApi, snapshot_download
from transformers import WhisperProcessor

# ---------- Config ----------
MODEL_IDS = {
    "small": "wolethereader/STORM-OS-ASR-SMALL-CT2",
    "large": "wolethereader/STORM-OS-ASR-LARGE-CT2",
}
VALID_LANGS = {"yo", "ha", "ig", "pcm", "en"}
ORG_NAME = "wolethereader"
MAX_CONCURRENT_INFERENCE = 2
MAX_QUEUE_SIZE = 20
RATE_LIMIT_REQUESTS = 30
RATE_LIMIT_WINDOW_SECONDS = 60
CHUNK_LENGTH_S = 20
STRIDE_LENGTH_S = 4
MAX_NEW_TOKENS = 225

# ---------- Logging ----------
logger = logging.getLogger("storm_os_api")
logger.setLevel(logging.INFO)
h = logging.StreamHandler(); h.setFormatter(logging.Formatter("%(message)s")); logger.addHandler(h)
def log_event(event, **fields):
    logger.info(json.dumps({"event": event, "ts": time.time(), **fields}))

# ---------- App + model loading (CTranslate2, resilient to missing repos) ----------
app = FastAPI(title="STORM-OS ASR API", version="2.0-ct2")
log_event("startup_begin")

HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    log_event("warning", message="HF_TOKEN not set - private model repos will fail to load")

MODELS, PROCESSORS = {}, {}
for size, repo in MODEL_IDS.items():
    try:
        log_event("loading_model", model_size=size, repo=repo)
        local_dir = snapshot_download(repo_id=repo, token=HF_TOKEN)
        PROCESSORS[size] = WhisperProcessor.from_pretrained(local_dir)
        MODELS[size] = ctranslate2.models.Whisper(
            local_dir, compute_type="int8", inter_threads=1, intra_threads=2
        )
        log_event("model_loaded_ok", model_size=size)
    except Exception as e:
        log_event("model_unavailable", model_size=size, repo=repo, error=str(e))
log_event("startup_complete", available_models=list(MODELS.keys()))

for size, model in MODELS.items():
    try:
        log_event("warmup_start", model_size=size)
        processor = PROCESSORS[size]
        silence = np.zeros(16000, dtype=np.float32)
        features = processor(silence, sampling_rate=16000, return_tensors="np").input_features
        features = ctranslate2.StorageView.from_array(features)
        prompt = processor.tokenizer.convert_tokens_to_ids(
            ["<|startoftranscript|>", "<|en|>", "<|transcribe|>", "<|notimestamps|>"]
        )
        model.generate(features, [prompt], max_length=5)
        log_event("warmup_complete", model_size=size)
    except Exception as e:
        log_event("warmup_failed", model_size=size, error=str(e))

try:
    from silero_vad import load_silero_vad
    VAD_MODEL = load_silero_vad()
    log_event("vad_loaded_ok")
except Exception as e:
    VAD_MODEL = None
    log_event("vad_load_failed", error=str(e))

# ---------- Auth: HF token AND must belong to the org ----------
security = HTTPBearer()
hf_api = HfApi()
_token_cache = {}
TOKEN_CACHE_TTL = 300

# Optional secondary access token for external collaborators who are not
# members of the HF org (e.g. contracted engineers). Set via the Space's
# secrets, never committed to source. To revoke access, delete this secret
# and restart the Space — does not affect normal org-token auth at all.
EXTERNAL_ACCESS_TOKEN = os.environ.get("EXTERNAL_ACCESS_TOKEN")

async def verify_org_token(creds: HTTPAuthorizationCredentials = Depends(security)):
    token = creds.credentials

    if EXTERNAL_ACCESS_TOKEN and token == EXTERNAL_ACCESS_TOKEN:
        log_event("auth_external_token_used")
        return "external-collaborator"

    now = time.time()
    cached = _token_cache.get(token)
    if cached and cached[1] > now:
        return cached[0]
    try:
        info = hf_api.whoami(token=token)
    except Exception:
        log_event("auth_failed_invalid_token")
        raise HTTPException(status_code=401, detail="Invalid or expired Hugging Face token")
    username = info.get("name", "unknown")
    user_orgs = [o.get("name") for o in info.get("orgs", [])]
    if ORG_NAME not in user_orgs:
        log_event("auth_failed_not_org_member", user=username, orgs=user_orgs)
        raise HTTPException(status_code=403, detail=f"Token does not belong to a member of '{ORG_NAME}'")
    _token_cache[token] = (username, now + TOKEN_CACHE_TTL)
    return username

# ---------- Rate limiting ----------
_rate_state = defaultdict(deque)
def check_rate_limit(username: str):
    now = time.time()
    window = _rate_state[username]
    while window and now - window[0] > RATE_LIMIT_WINDOW_SECONDS:
        window.popleft()
    if len(window) >= RATE_LIMIT_REQUESTS:
        raise HTTPException(status_code=429, detail="Rate limit exceeded, slow down")
    window.append(now)

# ---------- Bounded queue ----------
inference_semaphore = asyncio.Semaphore(MAX_CONCURRENT_INFERENCE)
_active_and_queued = 0
_queue_lock = asyncio.Lock()

async def acquire_queue_slot():
    global _active_and_queued
    async with _queue_lock:
        if _active_and_queued >= MAX_QUEUE_SIZE:
            raise HTTPException(status_code=429, detail="Server busy, queue full - retry shortly")
        _active_and_queued += 1

async def release_queue_slot():
    global _active_and_queued
    async with _queue_lock:
        _active_and_queued -= 1

# ---------- Audio helpers ----------
def load_and_normalize_audio(raw_bytes: bytes):
    try:
        audio, sr = sf.read(io.BytesIO(raw_bytes))
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"could not read audio file: {e}")
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    if sr != 16000:
        audio = librosa.resample(audio.astype("float32"), orig_sr=sr, target_sr=16000)
        sr = 16000
    return audio.astype("float32"), sr

def chunk_audio(audio, sr=16000, chunk_length_s=CHUNK_LENGTH_S, stride_length_s=STRIDE_LENGTH_S):
    chunk_len, stride_len = chunk_length_s * sr, stride_length_s * sr
    step = chunk_len - stride_len
    chunks = []
    for start in range(0, len(audio), step):
        chunk = audio[start:start + chunk_len]
        if len(chunk) == 0:
            break
        chunks.append(chunk)
        if start + chunk_len >= len(audio):
            break
    return chunks

def merge_chunk_texts(texts):
    if not texts:
        return ""
    merged = texts[0]
    for nxt in texts[1:]:
        mw, nw = merged.split(), nxt.split()
        best = 0
        for k in range(min(len(mw), len(nw), 10), 0, -1):
            if mw[-k:] == nw[:k]:
                best = k
                break
        merged = merged + " " + " ".join(nw[best:])
    return merged.strip()

def run_inference(audio, sr, model_size, lang_code):
    processor, model = PROCESSORS[model_size], MODELS[model_size]
    duration_s = len(audio) / sr
    chunks = [audio] if duration_s <= CHUNK_LENGTH_S else chunk_audio(audio, sr=sr)

    prompt = processor.tokenizer.convert_tokens_to_ids([
        "<|startoftranscript|>", f"<|{lang_code}|>", "<|transcribe|>", "<|notimestamps|>"
    ])

    feature_arrays = [
        processor(chunk, sampling_rate=sr, return_tensors="np").input_features[0]
        for chunk in chunks
    ]
    features = ctranslate2.StorageView.from_array(np.stack(feature_arrays))
    prompts = [prompt] * len(chunks)

    results = model.generate(
        features, prompts,
        max_length=MAX_NEW_TOKENS,
        no_repeat_ngram_size=3,
        repetition_penalty=1.3,
    )
    texts = [
        processor.tokenizer.decode(r.sequences_ids[0], skip_special_tokens=True).strip()
        for r in results
    ]

    return merge_chunk_texts(texts), duration_s, len(chunks)

# ---------- Endpoints ----------
@app.get("/")
def root():
    return {"status": "ok", "models_available": list(MODELS.keys()), "languages": sorted(VALID_LANGS),
            "engine": "ctranslate2-int8", "auth": f"Bearer <hf-token belonging to org '{ORG_NAME}'> required on /transcribe"}

@app.get("/health")
def health():
    return {"status": "healthy", "models_loaded": list(MODELS.keys())}

@app.post("/transcribe")
async def transcribe(
    request: Request,
    file: UploadFile = File(...),
    model_size: str = Form("small"),
    language: str = Form(...),
    username: str = Depends(verify_org_token),
):
    request_id = str(uuid.uuid4())
    t0 = time.time()
    check_rate_limit(username)

    if model_size not in MODELS:
        raise HTTPException(status_code=400, detail=f"model_size '{model_size}' not available. Loaded: {list(MODELS.keys())}")
    if language not in VALID_LANGS:
        raise HTTPException(status_code=400, detail=f"language must be one of {sorted(VALID_LANGS)}")

    raw = await file.read()
    audio, sr = load_and_normalize_audio(raw)

    await acquire_queue_slot()
    log_event("request_queued", request_id=request_id, user=username, model_size=model_size, language=language)
    try:
        async with inference_semaphore:
            log_event("inference_start", request_id=request_id, user=username)
            text, duration_s, n_chunks = await asyncio.to_thread(run_inference, audio, sr, model_size, language)
    finally:
        await release_queue_slot()

    elapsed = time.time() - t0
    log_event("request_complete", request_id=request_id, user=username, model_size=model_size,
               language=language, audio_duration_s=round(duration_s, 2), n_chunks=n_chunks, elapsed_s=round(elapsed, 2))

    return {"request_id": request_id, "model_size": model_size, "language": language, "text": text,
            "audio_duration_s": round(duration_s, 2), "chunks_processed": n_chunks, "elapsed_s": round(elapsed, 2)}

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    log_event("request_error", path=str(request.url.path), status=exc.status_code, detail=exc.detail)
    return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})