# ===========================
# main.py  (FastAPI + GROQ RAG) - Edited for conversation history & longer answers
# ===========================

import os
import uuid
import shutil
import tempfile
import traceback
import json
from pathlib import Path
from typing import List, Optional

from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware

import pdfplumber
import chromadb
from chromadb.config import Settings

from dotenv import load_dotenv
load_dotenv()

# ----------------------------
# ENV VARIABLES
# ----------------------------
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
EMBED_MODEL = os.getenv("EMBED_MODEL", "text-embedding-3-small")  # OpenAI embed name (or your embed model)
CHAT_MODEL = os.getenv("CHAT_MODEL", "llama-3.3-70b-versatile")
PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
STT_API_KEY = os.getenv("STT_API_KEY")  # Optional: separate STT API key (if not using OpenAI)
STT_PROVIDER = os.getenv("STT_PROVIDER", "openai")  # Options: "openai" (Whisper) or "assemblyai"
HISTORY_DIR = os.getenv("CHAT_HISTORY_DIR", "./chat_history")
os.makedirs(HISTORY_DIR, exist_ok=True)
# ----------------------------
# GROQ CLIENT
# ----------------------------
from groq import Groq
groq_client = Groq(api_key=GROQ_API_KEY)

# ----------------------------
# Embedding (OpenAI or SentenceTransformer)
# ----------------------------
USE_SENTENCE_TRANSFORMERS = False

try:
    from openai import OpenAI
    openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
except Exception:
    openai_client = None

try:
    if openai_client is None:
        from sentence_transformers import SentenceTransformer
        embedder = SentenceTransformer("all-MiniLM-L6-v2")
        USE_SENTENCE_TRANSFORMERS = True
except Exception:
    # If neither embedding provider available, raise a helpful message
    raise Exception("No embedding provider available (OpenAI or SentenceTransformers required). Please set OPENAI_API_KEY or install sentence-transformers.")

# ----------------------------
# Chroma DB
# ----------------------------
client = chromadb.Client(
    Settings(
        persist_directory=PERSIST_DIR,
        allow_reset=True
    )
)

# ----------------------------
# FastAPI App + CORS
# ----------------------------
app = FastAPI()

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

# ----------------------------
# Token / Word Chunking
# ----------------------------
try:
    import tiktoken
    enc = tiktoken.get_encoding("cl100k_base")
    TIKTOKEN_AVAILABLE = True
except Exception:
    TIKTOKEN_AVAILABLE = False


def split_text(text: str, chunk_tokens=800, overlap=200):
    """Split text using tokens if possible, else word-based."""
    if TIKTOKEN_AVAILABLE:
        tokens = enc.encode(text)
        chunks = []
        start = 0
        L = len(tokens)
        while start < L:
            end = min(start + chunk_tokens, L)
            chunks.append(enc.decode(tokens[start:end]))
            start = end - overlap if end - overlap > start else end
        return chunks
    else:
        words = text.split()
        chunk_size = int(chunk_tokens * 0.75)
        over = int(overlap * 0.75)
        chunks = []
        i = 0
        while i < len(words):
            chunks.append(" ".join(words[i:i + chunk_size]))
            i += chunk_size - over
        return chunks

def history_file(session_id: str) -> str:
    return os.path.join(HISTORY_DIR, f"{session_id}.json")


def load_history(session_id: str) -> list:
    path = history_file(session_id)
    if not os.path.exists(path):
        return []
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_message(session_id: str, role: str, text: str):
    history = load_history(session_id)
    history.append({
        "id": uuid.uuid4().hex,
        "role": role,
        "text": text
    })
    with open(history_file(session_id), "w", encoding="utf-8") as f:
        json.dump(history, f, ensure_ascii=False, indent=2)

# ----------------------------
# Embedding Function
# ----------------------------
def embed_texts(texts: List[str]) -> List[List[float]]:
    """Embed using OpenAI or sentence-transformers; always return native python floats."""
    if not USE_SENTENCE_TRANSFORMERS:
        # OpenAI embeddings
        resp = openai_client.embeddings.create(model=EMBED_MODEL, input=texts)
        data = getattr(resp, "data", resp.get("data", []))
        vectors = []
        for d in data:
            if isinstance(d, dict):
                vec = d.get("embedding")
            else:
                vec = getattr(d, "embedding", None)
            # convert to python floats
            vectors.append([float(x) for x in list(vec)])
        return vectors
    else:
        # SentenceTransformers returns numpy array; convert rows to python floats
        arr = embedder.encode(texts, normalize_embeddings=True)
        return [[float(x) for x in row] for row in arr]

# ----------------------------
# Extract text from PDF
# ----------------------------
def extract_pdf(path: str) -> List[dict]:
    pages = []
    with pdfplumber.open(path) as pdf:
        for i, p in enumerate(pdf.pages, start=1):
            text = p.extract_text() or ""
            if text.strip():
                pages.append({"page": i, "text": text.strip()})
    return pages

# ----------------------------
# Upload & Index
# ----------------------------
@app.post("/api/upload")
async def upload(files: List[UploadFile] = File(...), session_id: str = Form("default")):
    collection_name = f"session_{session_id}"
    chunks_all = []
    indexed_files = []

    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)

            for f in files:
                indexed_files.append(f.filename)

                saved = tmp / f.filename
                with saved.open("wb") as handle:
                    shutil.copyfileobj(f.file, handle)

                # Extract text
                pdf_pages = extract_pdf(str(saved))

                for pg in pdf_pages:
                    page_text = pg["text"]
                    page_num = pg["page"]

                    # Chunking
                    chunks = split_text(page_text)

                    for idx, ch in enumerate(chunks):
                        cid = f"{f.filename}__p{page_num}__c{idx}__{uuid.uuid4().hex[:8]}"
                        chunks_all.append({
                            "id": cid,
                            "text": ch,
                            "metadata": {"source": f.filename, "page": page_num, "chunk": idx}
                        })

        if not chunks_all:
            return JSONResponse({"success": False, "message": "No text extracted from uploaded files"}, status_code=400)

        # Embed in batches
        texts = [c["text"] for c in chunks_all]
        batch = 64
        vectors = []
        for i in range(0, len(texts), batch):
            vectors.extend(embed_texts(texts[i:i+batch]))

        # Convert any numpy / np.float32 values to native python floats
        def ensure_python_floats(vecs):
            """
            Convert vector-like objects to list[list[float]] using Python float types.
            Accepts: list of lists, numpy arrays, sentence-transformers arrays, etc.
            Returns: nested Python lists of native floats.
            """
            cleaned = []
            for v in vecs:
                # If v is a scalar vector-like (e.g., numpy array), convert to list
                try:
                    seq = list(v)
                except Exception:
                    seq = v
                # now ensure each element is python float
                cleaned.append([float(x) for x in seq])
            return cleaned

        vectors_clean = ensure_python_floats(vectors)

        # Upsert into Chroma using cleaned vectors
        collections = [c.name for c in client.list_collections()]
        col = client.get_collection(collection_name) if collection_name in collections else client.create_collection(collection_name)

        col.add(
            ids=[c["id"] for c in chunks_all],
            documents=texts,
            metadatas=[c["metadata"] for c in chunks_all],
            embeddings=vectors_clean
        )

        return {"success": True, "indexed_files": indexed_files}

    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"success": False, "message": str(e)}, status_code=500)

# ----------------------------
# Chat (RAG) - now supports 'history' (JSON string) to allow follow-ups
# ----------------------------
# Replace your existing /api/chat endpoint with this implementation
@app.post("/api/chat")
async def chat(
    session_id: str = Form(...),
    message: str = Form(...),
    top_k: int = Form(4),
    history: Optional[str] = Form(None),
):
    """
    Chat endpoint with:
     - RAG if session collection exists
     - Auto-indexing of long pasted user text when session not yet created
     - Direct GROQ fallback for greetings and short questions when no docs present
    """
    try:
        collection_name = f"session_{session_id}"
        collections = [c.name for c in client.list_collections()]

        # small helpers
        def is_greeting(text: str) -> bool:
            txt = text.strip().lower()
            greetings = ["hi", "hello", "hey", "good morning", "good afternoon", "good evening"]
            if txt in greetings:
                return True
            # short phrases containing greeting words
            if len(txt.split()) <= 3 and any(g in txt for g in greetings):
                return True
            return False

        def should_index_text(text: str) -> bool:
            t = text.strip()
            # heuristic: multiline or long => treat as notes to index
            if "\n" in t or len(t) >= 120 or len(t.split()) > 20:
                return True
            return False

        def run_direct_groq(prompt_text: str) -> str:
            resp = groq_client.chat.completions.create(
                model=CHAT_MODEL,
                messages=[
                    {"role": "system", "content": "You are an exam-focused AI tutor."},
                    {"role": "user", "content": prompt_text}
                ],
                temperature=0.1,
            )
            choice0 = (getattr(resp, "choices", None) or resp.get("choices", []))[0]
            if isinstance(choice0, dict):
                return choice0["message"]["content"].strip()
            else:
                return choice0.message.content.strip()

        # If collection exists -> normal RAG path
        if collection_name in collections:
            col = client.get_collection(collection_name)

            # embed query (ensure floats)
            q_vec = embed_texts([message])[0]
            q_vec = [float(x) for x in list(q_vec)]

            # retrieve
            results = col.query(query_embeddings=[q_vec], n_results=top_k, include=["metadatas", "documents", "distances"])
            docs = results.get("documents", [[]])[0]
            metadatas = results.get("metadatas", [[]])[0]

            context_blocks = []
            sources = []
            for md, doc in zip(metadatas, docs):
                src = md.get("source", "document")
                page = md.get("page")
                chunk = md.get("chunk")
                sources.append({"source": src, "page": page, "chunk": chunk})
                context_blocks.append(f"Source: {src} (page {page}, chunk {chunk})\n{doc}")

            context_text = "\n\n---\n\n".join(context_blocks) if context_blocks else ""

            # If no context found in retrieval, fallback to direct LLM (optionally)
            if not context_blocks:
                # you can change to "I don't know" if you want to strictly require documents
                answer = run_direct_groq(message)
                return {"answer": answer, "sources": []}

            # Build prompt for GROQ (RAG)
            system_prompt = """
                You are a professional study and knowledge assistant.
                
                Rules:
                1. Use uploaded documents as the PRIMARY source of truth.
                2. If the documents clearly contain the answer, respond strictly based on them.
                3. If the documents are weak, incomplete, or do NOT contain the answer:
                   - Answer confidently using your general knowledge.
                   - Do NOT say "I don't know".
                4. DO NOT mention personal names, phone numbers, emails, or identifiers.
                5. You MAY mention technologies, skills, tools, and project descriptions.
                6. If a resume is uploaded, refer to content in a generic way
                   (e.g., "the resume mentions Redis was used for caching").
                7. Keep answers concise, structured, and exam-focused.
                """

            messages_payload = [
                {"role": "system", "content": system_prompt},
                {"role": "system", "content": f"Context:\n{context_text}"},
                {"role": "user", "content": message}
            ]

            resp = groq_client.chat.completions.create(model=CHAT_MODEL, messages=messages_payload, temperature=0.0)
            choice0 = (getattr(resp, "choices", None) or resp.get("choices", []))[0]
            if isinstance(choice0, dict):
                answer = choice0["message"]["content"].strip()
            else:
                answer = choice0.message.content.strip()

            return {"answer": answer, "sources": sources}

        # If collection does NOT exist:
        # 1) greeting -> direct LLM
        if is_greeting(message):
            answer = run_direct_groq(message)
            return {"answer": answer, "sources": []}

        # 2) long pasted text -> auto-index and then run RAG against it
        if should_index_text(message):
            # chunk the message (recommended for long notes)
            chunks = split_text(message)
            if not chunks:
                chunks = [message]

            ids = []
            docs = []
            metadatas = []
            for idx, ch in enumerate(chunks):
                cid = f"user_text__{idx}__{uuid.uuid4().hex[:8]}"
                ids.append(cid)
                docs.append(ch)
                metadatas.append({"source": "user_input", "page": 1, "chunk": idx})

            # create collection
            col = client.create_collection(collection_name)

            # embed and ensure python floats
            batch = 64
            vecs = []
            for i in range(0, len(docs), batch):
                vecs.extend(embed_texts(docs[i:i+batch]))
            # convert / ensure floats
            vecs_clean = [[float(x) for x in list(v)] for v in vecs]

            col.add(ids=ids, documents=docs, metadatas=metadatas, embeddings=vecs_clean)

            # now run RAG on user query (same session)
            q_vec = embed_texts([message])[0]
            q_vec = [float(x) for x in list(q_vec)]
            results = col.query(query_embeddings=[q_vec], n_results=top_k, include=["metadatas", "documents"])
            docs = results.get("documents", [[]])[0]
            metadatas = results.get("metadatas", [[]])[0]
            context_blocks = []
            for md, doc in zip(metadatas, docs):
                src = md.get("source", "document")
                page = md.get("page")
                chunk = md.get("chunk")
                context_blocks.append(f"Source: {src} (page {page}, chunk {chunk})\n{doc}")
            context_text = "\n\n---\n\n".join(context_blocks)

            # call GROQ with context
            system_prompt = """
                You are a professional study and knowledge assistant.
                
                Rules:
                1. Use uploaded documents as the PRIMARY source of truth.
                2. If the documents clearly contain the answer, respond strictly based on them.
                3. If the documents are weak, incomplete, or do NOT contain the answer:
                   - Answer confidently using your general knowledge.
                   - Do NOT say "I don't know".
                4. DO NOT mention personal names, phone numbers, emails, or identifiers.
                5. You MAY mention technologies, skills, tools, and project descriptions.
                6. If a resume is uploaded, refer to content in a generic way
                   (e.g., "the resume mentions Redis was used for caching").
                7. Keep answers concise, structured, and exam-focused.
                """

            messages_payload = [
                {"role": "system", "content": system_prompt},
                {"role": "system", "content": f"Context:\n{context_text}"},
                {"role": "user", "content": message}
            ]
            resp = groq_client.chat.completions.create(model=CHAT_MODEL, messages=messages_payload, temperature=0.0)
            choice0 = (getattr(resp, "choices", None) or resp.get("choices", []))[0]
            if isinstance(choice0, dict):
                answer = choice0["message"]["content"].strip()
            else:
                answer = choice0.message.content.strip()
            save_message(session_id, "user", message)
            save_message(session_id, "ai", answer)
            return {"answer": answer, "sources": []}

        # 3) short question w/o docs -> direct GROQ LLM
        answer = run_direct_groq(message)
        save_message(session_id, "user", message)
        save_message(session_id, "ai", answer)
        return {"answer": answer, "sources": []}

    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"success": False, "message": str(e)}, status_code=500)

@app.get("/api/history/{session_id}")
async def get_history(session_id: str):
    try:
        return {
            "session_id": session_id,
            "messages": load_history(session_id)
        }
    except Exception as e:
        return JSONResponse(
            {"success": False, "message": str(e)},
            status_code=500
        )

@app.get("/api/sessions")
async def list_sessions():
    chats = []
    for fname in os.listdir(HISTORY_DIR):
        if fname.endswith(".json"):
            sid = fname.replace(".json", "")
            history = load_history(sid)
            if history:
                title = history[0]["text"][:40]
                chats.append({
                    "id": sid,
                    "title": title
                })
    return {"sessions": chats}

# ----------------------------
# Speech-to-Text (STT) Endpoint
# ----------------------------
import subprocess
import tempfile
import os
import shutil
import azure.cognitiveservices.speech as speechsdk
from fastapi import UploadFile, File, Form
from fastapi.responses import JSONResponse
import traceback

@app.post("/api/stt")
async def speech_to_text(
    audio: UploadFile = File(...),
    session_id: str = Form(...)
):
    try:
        speech_key = os.getenv("AZURE_SPEECH_KEY")
        speech_region = os.getenv("AZURE_SPEECH_REGION")

        if not speech_key or not speech_region:
            return JSONResponse(
                {"success": False, "message": "Azure Speech credentials not set"},
                status_code=500
            )

        # Save webm
        with tempfile.NamedTemporaryFile(delete=False, suffix=".webm") as tmp_webm:
            shutil.copyfileobj(audio.file, tmp_webm)
            webm_path = tmp_webm.name

        # Convert to wav (16kHz, mono, PCM)
        wav_path = webm_path.replace(".webm", ".wav")

        subprocess.run(
            [
                "ffmpeg", "-y",
                "-i", webm_path,
                "-ac", "1",
                "-ar", "16000",
                "-f", "wav",
                wav_path
            ],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )

        # Azure Speech
        speech_config = speechsdk.SpeechConfig(
            subscription=speech_key,
            region=speech_region
        )
        speech_config.speech_recognition_language = "en-US"

        audio_input = speechsdk.audio.AudioConfig(filename=wav_path)
        recognizer = speechsdk.SpeechRecognizer(
            speech_config=speech_config,
            audio_config=audio_input
        )

        result = recognizer.recognize_once()

        if result.reason == speechsdk.ResultReason.RecognizedSpeech:
            text = result.text.strip()
            return {
                "success": True,
                "text": text,
                "result": text,
                "transcription": text
            }

        elif result.reason == speechsdk.ResultReason.NoMatch:
            return JSONResponse(
                {"success": False, "message": "No speech detected"},
                status_code=400
            )

        else:
            return JSONResponse(
                {"success": False, "message": f"Azure STT failed: {result.reason}"},
                status_code=500
            )

    except Exception as e:
        traceback.print_exc()
        return JSONResponse(
            {"success": False, "message": f"STT failed: {str(e)}"},
            status_code=500
        )

    finally:
        for p in ["webm_path", "wav_path"]:
            try:
                os.remove(locals()[p])
            except:
                pass

# ----------------------------
# Uvicorn (if local)
# ----------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)