import os
from typing import List, Optional

from fastapi import FastAPI, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel

from pinecone import Pinecone
from sentence_transformers import SentenceTransformer
from openai import OpenAI


# ==============================
# 🔐 CONFIG (Use HF Secrets)
# ==============================

PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
SARVAM_API_KEY = os.getenv("SARVAM_API_KEY")
API_SECRET_KEY = os.getenv("API_SECRET_KEY")

PINECONE_INDEX = "prathamesh-portfolio"
EMBED_MODEL = "BAAI/bge-m3"
TOP_K = 5
SARVAM_MODEL = "sarvam-m"
SARVAM_BASE_URL = "https://api.sarvam.ai/v1"

# Validate secrets
if not PINECONE_API_KEY:
    raise ValueError("PINECONE_API_KEY not set in environment variables.")

if not SARVAM_API_KEY:
    raise ValueError("SARVAM_API_KEY not set in environment variables.")

if not API_SECRET_KEY:
    raise ValueError("API_SECRET_KEY must be set in environment variables.")


# ==============================
# 🚀 INITIALIZE COMPONENTS
# ==============================

pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX)

embedder = SentenceTransformer(EMBED_MODEL)

sarvam_client = OpenAI(
    api_key=SARVAM_API_KEY,
    base_url=SARVAM_BASE_URL
)


# ==============================
# 🏗 FASTAPI APP
# ==============================

app = FastAPI(
    title="Prathamesh Portfolio Chatbot",
    description="RAG-powered chatbot API with secret key authentication"
)

# Enable CORS (adjust origin in production)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

security = HTTPBearer()


# ==============================
# 📦 REQUEST MODEL
# ==============================

class ChatRequest(BaseModel):
    query: str
    tag_filter: Optional[List[str]] = None


# ==============================
# 🔎 RETRIEVAL
# ==============================

def retrieve(query: str, top_k: int = TOP_K, tag_filter: Optional[List[str]] = None):
    query_embedding = embedder.encode(
        query,
        normalize_embeddings=True
    ).tolist()

    pinecone_filter = None
    if tag_filter:
        pinecone_filter = {"tags": {"$in": tag_filter}}

    results = index.query(
        vector=query_embedding,
        top_k=top_k,
        include_metadata=True,
        filter=pinecone_filter
    )

    retrieved = []
    for match in results.get("matches", []):
        retrieved.append({
            "chunk_id": match["id"],
            "score": round(match["score"], 4),
            "text": match["metadata"]["text"],
            "tags": match["metadata"].get("tags", []),
        })

    return retrieved


# ==============================
# 🧠 PROMPT BUILDER
# ==============================

def build_prompt(query: str, retrieved_chunks: list) -> str:
    context_blocks = []

    for i, chunk in enumerate(retrieved_chunks, 1):
        context_blocks.append(
            f"[Context {i} | Relevance: {chunk['score']} | Type: {', '.join(chunk['tags'])}]\n{chunk['text']}"
        )

    context_str = "\n\n".join(context_blocks)

    return f"""
You are a helpful and professional AI assistant for Prathamesh Raut's portfolio chatbot.

Use ONLY the context provided below to answer.
If the answer isn't present, say:
"I don't have specific information about that in Prathamesh's profile."

Be concise and accurate.

---
CONTEXT:
{context_str}
---
QUESTION: {query}
ANSWER:
"""


# ==============================
# 🏷 TAG DETECTION
# ==============================

def detect_tags(query: str) -> Optional[List[str]]:
    query_lower = query.lower()

    tag_keywords = {
        "project": ["project", "built", "developed", "created", "worked on"],
        "skills": ["skill", "technologies", "tools", "tech stack"],
        "work_experience": ["work", "job", "company", "role", "position"],
        "education": ["college", "degree", "university", "cgpa", "gpa"],
        "publication": ["publish", "paper", "research", "ieee"],
        "certification": ["certified", "harvard", "databricks", "cisco"],
        "personal": ["hobbies", "interests", "language"],
    }

    detected = []

    for tag, keywords in tag_keywords.items():
        if any(keyword in query_lower for keyword in keywords):
            detected.append(tag)

    return detected if detected else None


# ==============================
# 🤖 ANSWER GENERATION
# ==============================

def answer(query: str, tag_filter: Optional[List[str]] = None) -> str:
    chunks = retrieve(query, tag_filter=tag_filter)

    if not chunks:
        return "I couldn't find relevant information in Prathamesh's knowledge base."

    prompt = build_prompt(query, chunks)

    response = sarvam_client.chat.completions.create(
        model=SARVAM_MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a professional portfolio assistant for Prathamesh Raut. "
                    "Answer only using the provided context."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        temperature=0.3,
        max_tokens=600,
    )

    return response.choices[0].message.content.strip()


# ==============================
# 🔐 AUTHENTICATED ENDPOINT
# ==============================

@app.post("/chat")
async def chat_endpoint(
    request: ChatRequest,
    credentials: HTTPAuthorizationCredentials = Depends(security),
):
    # Validate secret key
    if credentials.credentials != API_SECRET_KEY:
        raise HTTPException(status_code=401, detail="Invalid secret key.")

    try:
        tags = request.tag_filter or detect_tags(request.query)
        response = answer(request.query, tag_filter=tags)
        return {"answer": response}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ==============================
# ❤️ HEALTH CHECK
# ==============================

@app.get("/")
async def root():
    return {"message": "Prathamesh Portfolio Chatbot API is running!"}