import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama
from pymongo import MongoClient
from huggingface_hub import hf_hub_download

# --- 1. Database Setup ---
# Set this in your Hugging Face Space Secrets!
MONGO_URI = os.environ.get("MONGO_URI") 
try:
    client = MongoClient(MONGO_URI)
    db = client["legal_db"] # Replace with your database name
    collection = db["legal_records"] # Replace with your collection name
    print("Connected to MongoDB successfully.")
except Exception as e:
    print(f"MongoDB connection failed: {e}")

# --- 2. Model Setup ---
print("Downloading Llama-3.2 GGUF model...")
model_path = hf_hub_download(
    repo_id="bartowski/Llama-3.2-3B-Instruct-GGUF", 
    filename="Llama-3.2-3B-Instruct-Q4_K_M.gguf"
)

print("Loading LLM into memory...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048, 
    n_threads=2 
)
print("LLM loaded successfully.")

# --- 3. FastAPI App ---
app = FastAPI(title="Llama-3.2 Legal Generator API")

class GenerateRequest(BaseModel):
    query: str
    intent: str

@app.post("/generate")
def generate_response(request: GenerateRequest):
    # 1. Search MongoDB for relevant legal context
    # This searches your DB for records matching the intent
    db_results = list(collection.find(
        {"intent": request.intent}, 
        {"_id": 0, "answer": 1, "source": 1}
    ).limit(3)) # Limit to top 3 results to fit in context window
    
    # Format the DB results into a readable string
    context = ""
    sources = []
    if db_results:
        for idx, doc in enumerate(db_results):
            context += f"Fact {idx+1}: {doc.get('answer', '')}\n"
            if doc.get('source'):
                sources.append(doc['source'])
    else:
        context = "No specific legal precedent found in the database."
        sources.append("None")

    # 2. Build the Prompt for Llama
    
    prompt = f"""<|start_header_id|>system<|end_header_id|>
You are a legal AI assistant specializing in Indian women's rights. Use the provided Database Context to answer the user's question accurately. 
Database Context:
{context}<|eot_id|><|start_header_id|>user<|end_header_id|>
User Query: {request.query}
Intent: {request.intent}
Answer:<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""

    # 3. Generate the Answer
    output = llm(
        prompt,
        max_tokens=250,
        temperature=0.3,
        echo=False
    )
    
    return {
        "generated_answer": output['choices'][0]['text'].strip(),
        "db_context_used": context,
        "sources": list(set(sources)) # Return unique sources
    }