import os
import json
import torch
import logging
import traceback
from contextlib import asynccontextmanager

# ---------------------------------------------------------------------------
# Environment must be configured BEFORE importing transformers/huggingface_hub
# ---------------------------------------------------------------------------
os.environ.setdefault("HF_HOME", "/app/model_cache")
os.environ.setdefault("TRANSFORMERS_CACHE", "/app/model_cache")

_hf_token = os.getenv("HF_TOKEN")
if _hf_token:
    os.environ["HUGGING_FACE_HUB_TOKEN"] = _hf_token

from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from typing import List
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)
logging.getLogger("uvicorn.access").setLevel(logging.INFO)

MODEL_ID = os.getenv("MODEL_ID", "gsstec/LFM2-700M")

_num_threads = int(os.getenv("NUM_THREADS", os.cpu_count() or 4))
torch.set_num_threads(_num_threads)
torch.set_num_interop_threads(max(1, _num_threads // 2))

# BF16 needs AVX-512; HF CPU Spaces are typically AVX2 → use float32
_dtype = (
    torch.bfloat16
    if torch.backends.cpu.get_cpu_capability() >= "avx512"
    else torch.float32
)

model = None
tokenizer = None
_model_ready = False


# ---------------------------------------------------------------------------
# Lifespan — loads model after uvicorn binds the port
# ---------------------------------------------------------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, tokenizer, _model_ready
    _token_kwarg = {"token": _hf_token} if _hf_token else {}

    print(f"MODEL: {MODEL_ID}")
    print(f"dtype={_dtype}  threads={_num_threads}  authenticated={bool(_hf_token)}")

    try:
        print("Loading tokenizer...")
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_ID,
            clean_up_tokenization_spaces=False,
            **_token_kwarg,
        )
        tokenizer.padding_side = "left"
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        print("Tokenizer ready.")

        print("Loading model...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            dtype=_dtype,
            low_cpu_mem_usage=True,
            **_token_kwarg,
        )
        model.eval()
        _model_ready = True
        print("Model ready.")

    except Exception:
        traceback.print_exc()

    yield   # server is running


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------
app = FastAPI(title="LFM2-700M Inference API", lifespan=lifespan)


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    body = await request.body()
    safe_errors = [
        {k: (v.decode(errors="replace") if isinstance(v, bytes) else v)
         for k, v in err.items()}
        for err in exc.errors()
    ]
    return JSONResponse(
        status_code=422,
        content={"detail": safe_errors, "body_received": body.decode(errors="replace")},
    )


# ---------------------------------------------------------------------------
# Schemas
# ---------------------------------------------------------------------------
class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: List[Message]
    max_new_tokens: int = 256
    temperature: float = 0.3
    min_p: float = 0.15
    repetition_penalty: float = 1.05
    stream: bool = True


class ChatResponse(BaseModel):
    response: str


class AskRequest(BaseModel):
    question: str
    max_new_tokens: int = 128
    temperature: float = 0.3
    min_p: float = 0.15
    repetition_penalty: float = 1.05


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _check_ready():
    if not _model_ready:
        raise HTTPException(status_code=503, detail="Model is still loading, please retry shortly.")


def build_input_ids(messages: List[Message]) -> torch.Tensor:
    chat = [{"role": m.role, "content": m.content} for m in messages]
    encoding = tokenizer.apply_chat_template(
        chat,
        add_generation_prompt=True,
        return_tensors="pt",
        tokenize=True,
        return_dict=True,
    )
    return encoding["input_ids"].to(model.device)


def _generate(input_ids: torch.Tensor, req, streamer=None) -> torch.Tensor:
    kwargs = dict(
        input_ids=input_ids,
        do_sample=True,
        temperature=req.temperature,
        min_p=req.min_p,
        repetition_penalty=req.repetition_penalty,
        max_new_tokens=req.max_new_tokens,
        use_cache=True,
    )
    if streamer:
        kwargs["streamer"] = streamer
    with torch.inference_mode():
        return model.generate(**kwargs)


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------
@app.get("/")
def root():
    return {"status": "ok", "model": MODEL_ID, "ready": _model_ready}


@app.get("/health")
def health():
    return {
        "status": "ok" if _model_ready else "loading",
        "model": MODEL_ID,
        "ready": _model_ready,
        "dtype": str(_dtype),
        "cpu_threads": _num_threads,
        "authenticated": bool(_hf_token),
    }


@app.post("/chat")
async def chat(request: Request):
    _check_ready()
    try:
        body = await request.body()
        req = ChatRequest(**json.loads(body))
    except Exception as e:
        raise HTTPException(status_code=422, detail=f"Invalid request body: {e}")

    input_ids = build_input_ids(req.messages)

    if req.stream:
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        Thread(target=_generate, args=(input_ids, req, streamer), daemon=True).start()
        return StreamingResponse((tok for tok in streamer), media_type="text/plain")

    output = _generate(input_ids, req)
    new_tokens = output[0][input_ids.shape[-1]:]
    return ChatResponse(response=tokenizer.decode(new_tokens, skip_special_tokens=True))


@app.post("/ask")
async def ask(request: Request):
    _check_ready()
    try:
        body = await request.body()
        req = AskRequest(**json.loads(body))
    except Exception as e:
        raise HTTPException(status_code=422, detail=f"Invalid request body: {e}")

    input_ids = build_input_ids([Message(role="user", content=req.question)])
    output = _generate(input_ids, req)
    new_tokens = output[0][input_ids.shape[-1]:]
    return ChatResponse(response=tokenizer.decode(new_tokens, skip_special_tokens=True))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="warning", access_log=True)