import os
import shutil
import logging
from typing import List, Dict, Any
from contextlib import asynccontextmanager

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import chromadb
from chromadb.utils import embedding_functions
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from llama_cpp import Llama
import uvicorn

# ------------------------------------------------------------
# Configuration
# ------------------------------------------------------------
MODEL_PATH = os.getenv("MODEL_PATH", "/home/orangepi/daemon/orangepi_service/models/qwen2.5-0.5b-instruct-q8_0.gguf")
N_CTX = int(os.getenv("N_CTX", "1024"))          # Smaller context for Orange Pi
N_THREADS = int(os.getenv("N_THREADS", "4"))
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "rag_docs")
CHROMA_DIR = os.getenv("CHROMA_DIR", "/home/orangepi/daemon/orangepi_service/models/chroma_db")
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "/home/orangepi/daemon/orangepi_service/models/uploaded_pdfs")
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "500"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "50"))
TOP_K = int(os.getenv("TOP_K", "3"))             # Number of chunks to retrieve

# Generation defaults
TEMP = float(os.getenv("TEMP", "0.3"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256"))

# ------------------------------------------------------------
# Logging
# ------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("orangepi-rag")

# ------------------------------------------------------------
# Global objects (initialised at startup)
# ------------------------------------------------------------
llm = None
chroma_client = None
collection = None
embedding_fn = None

# ------------------------------------------------------------
# Initialisation
# ------------------------------------------------------------
def init_directories():
    os.makedirs(UPLOAD_DIR, exist_ok=True)
    os.makedirs(CHROMA_DIR, exist_ok=True)

def init_chroma():
    global chroma_client, collection, embedding_fn
    chroma_client = chromadb.PersistentClient(path=CHROMA_DIR)
    # Use sentence-transformers embedding function (runs on CPU)
    embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="all-MiniLM-L6-v2"
    )
    # Get or create collection
    collection = chroma_client.get_or_create_collection(
        name=COLLECTION_NAME,
        embedding_function=embedding_fn
    )
    logger.info(f"Chroma initialised, collection '{COLLECTION_NAME}' has {collection.count()} documents")

def init_llm():
    global llm
    logger.info(f"Loading GGUF model from {MODEL_PATH} ...")
    llm = Llama(
        model_path=MODEL_PATH,
        n_ctx=N_CTX,
        n_threads=N_THREADS,
        verbose=False,
    )
    logger.info("LLM loaded.")

def index_pdf(file_path: str, doc_id: str) -> int:
    """Load PDF, split, embed, and add to Chroma. Returns number of chunks."""
    loader = PyPDFLoader(file_path)
    documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", " ", ""]
    )
    chunks = text_splitter.split_documents(documents)

    if not chunks:
        logger.warning(f"No text extracted from {file_path}")
        return 0

    ids = [f"{doc_id}_{i}" for i in range(len(chunks))]
    texts = [chunk.page_content for chunk in chunks]
    metadatas = [{"source": doc_id, "page": chunk.metadata.get("page", 0)} for chunk in chunks]

    collection.add(
        ids=ids,
        documents=texts,
        metadatas=metadatas
    )
    logger.info(f"Indexed {len(chunks)} chunks from {doc_id}")
    return len(chunks)

def retrieve_relevant_chunks(question: str) -> List[str]:
    """Return top-K document chunks as strings."""
    results = collection.query(query_texts=[question], n_results=TOP_K)
    if results and results['documents']:
        return results['documents'][0]
    return []

def generate_answer(question: str, context_chunks: List[str]) -> str:
    """Generate answer using local LLM with retrieved context."""
    if not context_chunks:
        prompt = f"""Question: {question}\n\nNo relevant documents found. Answer based on your general knowledge.\nAnswer:"""
    else:
        context = "\n\n".join(context_chunks)
        prompt = f"""Use only the following context to answer the question. If the answer is not in the context, say "I don't know".

Context:
{context}

Question: {question}

Answer:"""
    output = llm.create_completion(
        prompt="Hello how are you?",
        max_tokens=MAX_TOKENS,
        temperature=TEMP,
        stop=None,
        echo=False
    )
    print(output)
    return output["choices"][0]["text"].strip()

# ------------------------------------------------------------
# FastAPI lifespan
# ------------------------------------------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    init_directories()
    init_chroma()
    init_llm()
    yield
    # Optional: cleanup
    logger.info("Shutting down OrangePi RAG service.")

app = FastAPI(
    title="OrangePi RAG Service",
    description="PDF upload, retrieval, and answer generation using a local GGUF model.",
    version="1.0",
    lifespan=lifespan,
)

# ------------------------------------------------------------
# Request/Response models
# ------------------------------------------------------------
class QuestionRequest(BaseModel):
    question: str

class AnswerResponse(BaseModel):
    answer: str
    retrieved_chunks: List[str] = []   # for debugging

class DocumentInfo(BaseModel):
    doc_id: str
    filename: str
    chunks: int

# ------------------------------------------------------------
# API Endpoints
# ------------------------------------------------------------
@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": llm is not None, "docs_indexed": collection.count() if collection else 0}

@app.post("/upload")
async def upload_pdf(file: UploadFile = File(...)):
    """Upload a PDF file, store it, and index its contents."""
    if not file.filename.endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed")

    # Generate a unique doc_id from filename + timestamp
    doc_id = file.filename.replace(' ', '_').replace('/', '_')
    base, ext = os.path.splitext(doc_id)
    doc_id = f"{base}_{os.urandom(4).hex()}{ext}"
    file_path = os.path.join(UPLOAD_DIR, doc_id)

    try:
        # Save file
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        # Index the PDF
        num_chunks = index_pdf(file_path, doc_id)

        return JSONResponse(content={
            "message": "PDF uploaded and indexed",
            "doc_id": doc_id,
            "original_filename": file.filename,
            "chunks_indexed": num_chunks
        })
    except Exception as e:
        logger.exception("Upload/index failed")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/documents")
async def list_documents():
    """List all indexed documents with metadata."""
    # Chroma doesn't store document-level metadata easily; we maintain a separate index.
    # For simplicity, we'll read from the collection's metadata (if we stored source per chunk)
    # Actually we stored 'source' in each chunk's metadata. Let's aggregate.
    all_metas = collection.get()  # returns ids, metadatas, documents
    if not all_metas or not all_metas['metadatas']:
        return {"documents": []}

    doc_map = {}
    for meta in all_metas['metadatas']:
        source = meta.get('source')
        if source:
            if source not in doc_map:
                doc_map[source] = 0
            doc_map[source] += 1
    docs = [{"doc_id": doc_id, "chunks": count} for doc_id, count in doc_map.items()]
    return {"documents": docs}

@app.post("/answer", response_model=AnswerResponse)
async def answer_question(req: QuestionRequest):
    """RAG answer: retrieve relevant chunks and generate answer."""
    if llm is None:
        raise HTTPException(status_code=503, detail="LLM not loaded")

    # Retrieve
    chunks = retrieve_relevant_chunks(req.question)
    # Generate
    answer = generate_answer(req.question, chunks)

    return AnswerResponse(answer=answer, retrieved_chunks=chunks)

@app.delete("/documents/{doc_id}")
async def delete_document(doc_id: str):
    """
    Delete a document and its chunks from the vector database.
    Optionally also delete the uploaded PDF file.
    """
    try:
        # 1. Retrieve all chunks with metadata source == doc_id
        # Chroma doesn't support direct delete by metadata, so we get all chunk ids with that source
        all_chunks = collection.get(where={"source": doc_id})
        if not all_chunks or not all_chunks['ids']:
            raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")

        # 2. Delete the chunks
        collection.delete(ids=all_chunks['ids'])

        # 3. (Optional) Delete the physical PDF file if it exists
        pdf_path = os.path.join(UPLOAD_DIR, doc_id)
        if os.path.exists(pdf_path):
            os.remove(pdf_path)

        return JSONResponse(content={"message": f"Document {doc_id} deleted", "chunks_deleted": len(all_chunks['ids'])})
    except Exception as e:
        logger.exception(f"Failed to delete document {doc_id}")
        raise HTTPException(status_code=500, detail=str(e))

# ------------------------------------------------------------
# Run
# ------------------------------------------------------------
if __name__ == "__main__":
    uvicorn.run(
        "service:app",
        host="0.0.0.0",
        port=8080,
        reload=False,
        log_level="info"
    )