from fastapi import FastAPI, Header, HTTPException, Body
from sentence_transformers import SentenceTransformer
import uvicorn
import os
from pathlib import Path
import chromadb
from chromadb.config import Settings
from chromadb.server.fastapi import FastAPI as ChromaFastAPI
import torch
import logging
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import List

# --- Logging Setup ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(process)d - %(levelname)s - %(message)s",
)
logger = logging.getLogger("HarrierService")

# --- CPU Optimization ---
# Crucial for multi-worker setups: prevents workers from fighting over the same CPU cores
torch.set_num_threads(1)

app = FastAPI(title="Harrier OSS 0.6B Embedder Production Service")

# --- Configuration ---
MODEL_NAME = "microsoft/harrier-oss-v1-0.6b"
API_KEY_REQUIRED = os.getenv("API_KEY", "Azure123")
QUERY_INSTRUCTION = "Instruct: Retrieve relevant passages that answer the query\nQuery: "
CHROMA_HOST = os.getenv("CHROMA_HOST", "0.0.0.0")
CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY") or "/data/chroma_data"

# Force CPU since it's a CPU-only server
device = "cpu"
logger.info(f"[*] Starting service worker on device: {device}")

# Internal thread pool for each worker to handle the handoff
executor = ThreadPoolExecutor(max_workers=1)

# Model loaded globally per worker process
model = None


def ensure_bucket_path(path: str) -> str:
    bucket_path = Path(path)
    try:
        bucket_path.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        raise RuntimeError(
            f"Storage path is not writable: {bucket_path}. "
            "This server must run on a Hugging Face Space with a mounted bucket."
        ) from exc
    if not os.access(bucket_path, os.W_OK):
        raise RuntimeError(
            f"Storage path is not writable: {bucket_path}. "
            "This server must run on a Hugging Face Space with a mounted bucket."
        )
    return str(bucket_path)


chroma_persist_directory = ensure_bucket_path(CHROMA_PERSIST_DIRECTORY)
chroma_settings = Settings(persist_directory=chroma_persist_directory)
chroma_server = ChromaFastAPI(chroma_settings)
app.mount("/chroma", chroma_server.app())

def write_bucket_probe(path: str) -> None:
    probe_path = Path(path) / "test.md"
    probe_path.write_text(
        "Chroma bucket probe: if you can read this file, the Space can write to /data.\n",
        encoding="utf-8",
    )


def seed_chroma_data(path: str) -> None:
    client = chromadb.PersistentClient(path=path)
    collection = client.get_or_create_collection(name="knowledge_base")
    if collection.count() > 0:
        logger.info("[*] Chroma already has data; skipping seed.")
        return

    documents = [
        "Chroma is a lightweight, open-source vector database built for AI.",
        "Python is a high-level programming language used extensively in data science.",
        "The celestial body closest to Earth is the Moon.",
    ]
    metadatas = [
        {"category": "tech", "source": "docs"},
        {"category": "tech", "source": "wiki"},
        {"category": "science", "source": "space-facts"},
    ]
    ids = ["doc1", "doc2", "doc3"]
    collection.add(documents=documents, metadatas=metadatas, ids=ids)
    logger.info("[+] Seeded Chroma with sample documents.")


@app.on_event("startup")
def load_model():
    global model
    write_bucket_probe(chroma_persist_directory)
    seed_chroma_data(chroma_persist_directory)
    logger.info(f"[*] Loading Harrier OSS 0.6B model: {MODEL_NAME}...")
    try:
        model = SentenceTransformer(MODEL_NAME, trust_remote_code=True, device=device)
        logger.info("[+] Model loaded successfully into worker.")
    except Exception as e:
        logger.error(f"[-] Failed to load model: {e}")
        raise e


def compute_embeddings(processed_input):
    """Heavy computation logic."""
    return model.encode(processed_input, show_progress_bar=False).tolist()


@app.post("/embed")
async def get_embeddings(
    input: List[str] = Body(...),
    is_query: bool = Body(False),
    x_api_key: str = Header(None),
):
    # 1. API Key Validation
    if x_api_key != API_KEY_REQUIRED:
        raise HTTPException(status_code=403, detail="Invalid API Key")

    # 2. Input Validation
    if not input:
        raise HTTPException(status_code=400, detail="Input list cannot be empty")

    # 3. Apply Harrier-specific instruction
    processed_input = [f"{QUERY_INSTRUCTION}{text}" for text in input] if is_query else input

    # 4. Offload to ThreadPool
    loop = asyncio.get_event_loop()
    try:
        logger.info(f"[*] Processing {len(input)} items (is_query={is_query})")
        embeddings = await loop.run_in_executor(executor, compute_embeddings, processed_input)
        return {"embeddings": embeddings}
    except Exception as e:
        logger.error(f"Error during embedding: {e}")
        raise HTTPException(status_code=500, detail="Internal embedding error")


@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "model": MODEL_NAME,
        "device": device,
        "worker_pid": os.getpid(),
    }


if __name__ == "__main__":
    # --- Deployment Config for 16GB RAM / CPU ---
    # With a 0.6B model (~2.5GB RAM), 4 workers = ~10GB.
    # This leaves 6GB for OS and overhead, which is safe and handles concurrency well.
    uvicorn.run(
        "embedder_service:app",
        host="0.0.0.0",
        workers=4,  # This creates 4 independent processes
        access_log=False,  # Disable for slightly better performance
    )