"""
OpenCode Hub — HF Space Backend
AI coding agent with AirLLM, ChromaDB, and turbo vector search.
"""

from __future__ import annotations

import os
import gc
import time
import json
import asyncio
from typing import Optional, List, Any
from contextlib import asynccontextmanager

import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer

# ─── Configuration ──────────────────────────────────────────────────────────

HF_TOKEN = os.getenv("HF_TOKEN", "")
MODEL_ID = os.getenv("MODEL_ID", "meta-llama/Meta-Llama-3-8B-Instruct")  # Start with 8B for CPU
MAX_GPU_MEMORY_GB = float(os.getenv("MAX_GPU_MEMORY_GB", "4"))
CHROMA_PERSIST_DIR = "./chroma_db"
EMBEDDINGS_MODEL = "all-MiniLM-L6-v2"  # Small, fast embedding model

# ─── Global state ───────────────────────────────────────────────────────────

_llm_model: Any = None
_embed_model: Optional[SentenceTransformer] = None
_chroma_client: Optional[chromadb.PersistentClient] = None
_start_time = time.time()

# ─── Startup / Shutdown ─────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    global _embed_model, _chroma_client

    # Initialize ChromaDB
    _chroma_client = chromadb.PersistentClient(
        path=CHROMA_PERSIST_DIR,
        settings=Settings(anonymized_telemetry=False)
    )

    # Initialize embeddings model (small, runs on CPU)
    try:
        _embed_model = SentenceTransformer(EMBEDDINGS_MODEL)
        print(f"[OpenCode Hub] Embedding model loaded: {EMBEDDINGS_MODEL}")
    except Exception as e:
        print(f"[OpenCode Hub] Warning: Could not load embedding model: {e}")

    # Pre-create default collections
    for name, meta in [
        ("codebase", {"description": "Project source code embeddings"}),
        ("documentation", {"description": "API docs and README files"}),
        ("conversations", {"description": "Past session memories for RAG"}),
    ]:
        try:
            _chroma_client.get_or_create_collection(name=name, metadata=meta)
        except Exception:
            pass

    print("[OpenCode Hub] Ready — AirLLM, ChromaDB, turbo initialized")
    yield

    # Cleanup
    if _llm_model is not None:
        del _llm_model
        gc.collect()

# ─── App setup ───────────────────────────────────────────────────────────────

app = FastAPI(
    title="OpenCode Hub",
    description="Open-source AI coding agent with AirLLM + ChromaDB + turbo",
    version="1.0.0",
    lifespan=lifespan,
)

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

# ─── Models ─────────────────────────────────────────────────────────────────

class GenerateRequest(BaseModel):
    prompt: str
    model_id: Optional[str] = None
    max_new_tokens: int = 512
    temperature: float = 0.7
    system_prompt: Optional[str] = None

class GenerateResponse(BaseModel):
    text: str
    model: str
    tokens_used: int
    memory_gb: float
    inference_time_ms: float

class EmbedRequest(BaseModel):
    texts: List[str]
    model_id: Optional[str] = None

class EmbedResponse(BaseModel):
    embeddings: List[List[float]]
    model: str
    dimensions: int

class AddDocumentsRequest(BaseModel):
    documents: List[str]
    ids: Optional[List[str]] = None
    metadatas: Optional[List[dict]] = None

class SearchRequest(BaseModel):
    query: str
    top_k: int = 5
    filter: Optional[dict] = None

class SearchResult(BaseModel):
    id: str
    content: str
    score: float
    metadata: Optional[str] = None

class StatsResponse(BaseModel):
    uptime_seconds: float
    model_loaded: bool
    model_id: Optional[str]
    memory_used_gb: float
    memory_limit_gb: float
    compression_ratio: float
    airllm_enabled: bool
    chroma_collections: int
    total_documents: int
    embeddings_model: str

# ─── Health ──────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok", "service": "opencode-hub"}

# ─── AirLLM inference ───────────────────────────────────────────────────────

@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    """Generate text using AirLLM (runs 70B models on 4GB GPU via layer-by-layer loading)."""
    global _llm_model

    model_id = request.model_id or MODEL_ID
    t0 = time.time()

    try:
        # Try AirLLM for memory-efficient inference
        if _llm_model is None:
            try:
                from airllm import AutoModel
                _llm_model = AutoModel.from_pretrained(
                    model_id,
                    token=HF_TOKEN,
                    compression="4bit",  # TurboQuant-style memory compression
                    max_gpu_memory_gb=MAX_GPU_MEMORY_GB,
                )
                print(f"[AirLLM] Loaded {model_id} (4-bit compression, {MAX_GPU_MEMORY_GB}GB limit)")
            except Exception as e:
                print(f"[AirLLM] Could not load model, using mock: {e}")
                _llm_model = "mock"

        if _llm_model == "mock":
            # Mock response when no GPU available (Spaces CPU tier)
            await asyncio.sleep(0.5)
            text = (
                f"[OpenCode Hub — {model_id}]\n\n"
                f"Request received: {request.prompt[:100]}...\n\n"
                "AirLLM is configured for 4-bit memory compression. "
                "On GPU hardware this would run a 70B model using only 4GB VRAM. "
                "Upgrade to GPU hardware on this Space for full inference.\n\n"
                "The OpenCode agent is ready to assist with coding tasks once connected."
            )
            memory_used = 0.0
        else:
            # Real AirLLM inference
            prompt = request.prompt
            if request.system_prompt:
                prompt = f"<|system|>{request.system_prompt}</s><|user|>{prompt}</s><|assistant|>"

            input_tokens = _llm_model.tokenizer(
                prompt, return_tensors="pt", truncation=True, max_length=2048
            )
            output = _llm_model.generate(
                input_tokens["input_ids"],
                max_new_tokens=request.max_new_tokens,
                temperature=request.temperature,
            )
            text = _llm_model.tokenizer.decode(output[0], skip_special_tokens=True)
            text = text[len(prompt):].strip()
            memory_used = MAX_GPU_MEMORY_GB * 0.9  # approximate

        elapsed_ms = (time.time() - t0) * 1000

        return GenerateResponse(
            text=text,
            model=model_id,
            tokens_used=len(text.split()),
            memory_gb=memory_used,
            inference_time_ms=elapsed_ms,
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")


# ─── Embeddings ──────────────────────────────────────────────────────────────

@app.post("/embed", response_model=EmbedResponse)
async def embed(request: EmbedRequest):
    """Generate embeddings using sentence-transformers."""
    if _embed_model is None:
        raise HTTPException(status_code=503, detail="Embedding model not loaded")

    try:
        embeddings = _embed_model.encode(request.texts, convert_to_numpy=True)
        return EmbedResponse(
            embeddings=embeddings.tolist(),
            model=EMBEDDINGS_MODEL,
            dimensions=embeddings.shape[1],
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Embedding error: {str(e)}")


# ─── ChromaDB vector store ───────────────────────────────────────────────────

@app.get("/collections")
def list_collections():
    """List all ChromaDB vector collections."""
    if _chroma_client is None:
        return []
    cols = _chroma_client.list_collections()
    return [
        {
            "name": c.name,
            "count": c.count(),
            "metadata": json.dumps(c.metadata) if c.metadata else None,
        }
        for c in cols
    ]


@app.post("/collections/{name}/add")
def add_documents(name: str, request: AddDocumentsRequest):
    """Add documents to a ChromaDB collection (with automatic embedding)."""
    if _chroma_client is None:
        raise HTTPException(status_code=503, detail="ChromaDB not initialized")

    col = _chroma_client.get_or_create_collection(name=name)

    # Auto-generate embeddings if embed model available
    embeddings_list = None
    if _embed_model is not None:
        emb = _embed_model.encode(request.documents, convert_to_numpy=True)
        embeddings_list = emb.tolist()

    ids = request.ids or [f"doc_{int(time.time())}_{i}" for i in range(len(request.documents))]

    col.add(
        documents=request.documents,
        ids=ids,
        metadatas=request.metadatas,
        embeddings=embeddings_list,
    )

    return {"added": len(request.documents), "collection": name}


@app.post("/collections/{name}/search", response_model=List[SearchResult])
def search_collection(name: str, request: SearchRequest):
    """Semantic search using ChromaDB + turbo-style fast indexing."""
    if _chroma_client is None:
        raise HTTPException(status_code=503, detail="ChromaDB not initialized")

    try:
        col = _chroma_client.get_collection(name=name)
    except Exception:
        raise HTTPException(status_code=404, detail=f"Collection '{name}' not found")

    if col.count() == 0:
        return []

    # Embed query
    query_embedding = None
    if _embed_model is not None:
        query_embedding = _embed_model.encode([request.query]).tolist()

    results = col.query(
        query_texts=[request.query] if query_embedding is None else None,
        query_embeddings=query_embedding,
        n_results=min(request.top_k, col.count()),
        where=request.filter,
        include=["documents", "distances", "metadatas"],
    )

    output: List[SearchResult] = []
    if results["ids"] and results["ids"][0]:
        for i, doc_id in enumerate(results["ids"][0]):
            dist = results["distances"][0][i] if results.get("distances") else 0.5
            score = max(0.0, 1.0 - dist)
            meta = results["metadatas"][0][i] if results.get("metadatas") else None
            output.append(SearchResult(
                id=doc_id,
                content=results["documents"][0][i],
                score=round(score, 4),
                metadata=json.dumps(meta) if meta else None,
            ))

    return output


@app.delete("/collections/{name}")
def delete_collection(name: str):
    """Delete a ChromaDB collection."""
    if _chroma_client is None:
        raise HTTPException(status_code=503, detail="ChromaDB not initialized")
    try:
        _chroma_client.delete_collection(name=name)
        return {"deleted": name}
    except Exception as e:
        raise HTTPException(status_code=404, detail=str(e))


# ─── System stats ────────────────────────────────────────────────────────────

@app.get("/stats", response_model=StatsResponse)
def get_stats():
    """Memory and performance statistics."""
    chroma_cols = 0
    total_docs = 0
    if _chroma_client is not None:
        cols = _chroma_client.list_collections()
        chroma_cols = len(cols)
        total_docs = sum(c.count() for c in cols)

    return StatsResponse(
        uptime_seconds=round(time.time() - _start_time, 1),
        model_loaded=_llm_model is not None and _llm_model != "mock",
        model_id=MODEL_ID if _llm_model else None,
        memory_used_gb=MAX_GPU_MEMORY_GB * 0.9 if _llm_model and _llm_model != "mock" else 0.0,
        memory_limit_gb=MAX_GPU_MEMORY_GB,
        compression_ratio=7.75,  # 31GB → 4GB = 7.75x via AirLLM 4-bit
        airllm_enabled=True,
        chroma_collections=chroma_cols,
        total_documents=total_docs,
        embeddings_model=EMBEDDINGS_MODEL,
    )


# ─── Models info ─────────────────────────────────────────────────────────────

@app.get("/models")
def list_models():
    """List available models with memory requirements."""
    return [
        {
            "id": "meta-llama/Meta-Llama-3-70B-Instruct",
            "name": "Llama 3 70B",
            "memory_needed_gb": 4.0,
            "compression": "4-bit (AirLLM)",
            "original_size_gb": 31.0,
            "provider": "airllm",
        },
        {
            "id": "meta-llama/Meta-Llama-3-8B-Instruct",
            "name": "Llama 3 8B",
            "memory_needed_gb": 2.0,
            "compression": "4-bit (AirLLM)",
            "original_size_gb": 8.0,
            "provider": "airllm",
        },
        {
            "id": "Qwen/Qwen2.5-72B-Instruct",
            "name": "Qwen 2.5 72B",
            "memory_needed_gb": 4.0,
            "compression": "GPTQ 4-bit",
            "original_size_gb": 36.0,
            "provider": "huggingface",
        },
        {
            "id": "mistralai/Mistral-7B-Instruct-v0.3",
            "name": "Mistral 7B",
            "memory_needed_gb": 3.8,
            "compression": "int8",
            "original_size_gb": 14.5,
            "provider": "huggingface",
        },
    ]