import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers"
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp/sentence_transformers"
os.environ["TORCH_HOME"] = "/tmp/torch"

import json
from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel
from typing import Optional
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer

# Import from autonomous agent
from agent_langchain import (
    process_with_agent,
    get_conversation_history,
    classify_ticket,
    call_routing,
    get_kb_collection,
    encoder,
    conversations
)

app = FastAPI(title="Smart Helpdesk AI Agent - Autonomous")

# Request Models
class TicketRequest(BaseModel):
    text: str
    conversation_id: Optional[str] = None
    user_email: Optional[str] = None

# Persistent Chroma settings
CHROMA_PATH = "/tmp/chroma"
COLLECTION_NAME = "knowledge_base"

# -------------------------------
# KB Setup Endpoint
# -------------------------------
@app.post("/setup")
async def setup_kb(kb_file: UploadFile = File(...)):
    """Upload and index knowledge base."""
    try:
        content_bytes = await kb_file.read()
        data = json.loads(content_bytes)

        if not isinstance(data, list):
            raise HTTPException(status_code=400, detail="JSON must be a list of items.")

        print(f"📘 Loaded {len(data)} items from {kb_file.filename}")

        chroma_client = chromadb.PersistentClient(
            path=CHROMA_PATH,
            settings=Settings(anonymized_telemetry=False, allow_reset=True)
        )
        collection = chroma_client.get_or_create_collection(COLLECTION_NAME)

        if collection.count() > 0:
            print(f"🧹 Clearing {collection.count()} existing records...")
            collection.delete(ids=collection.get()['ids'])

        texts, ids, metadatas = [], [], []
        for i, item in enumerate(data):
            text = item.get("answer") or item.get("text") or item.get("content") or ""
            item_id = item.get("id") or str(i)
            category = item.get("category", "")
            
            if not text:
                print(f"⚠️ Skipping item {i} - no text content")
                continue
            
            combined_text = f"Category: {category}. {text}" if category else text
            texts.append(combined_text)
            ids.append(str(item_id))
            metadatas.append({"id": str(item_id), "category": category, "original_index": i})

        if not texts:
            raise HTTPException(status_code=400, detail="No valid text content found in JSON.")

        print("🧠 Generating embeddings...")
        embeddings = encoder.encode(texts, show_progress_bar=True).tolist()

        print("💾 Adding to ChromaDB...")
        collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas)

        # Update global reference
        import agent_langchain
        agent_langchain.kb_collection = collection

        print(f"✅ Successfully added {collection.count()} records")
        return {"message": "Knowledge base initialized", "count": collection.count()}

    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON file.")
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Setup failed: {str(e)}")

# -------------------------------
# MAIN ORCHESTRATE ENDPOINT (Autonomous Agent)
# -------------------------------
@app.post("/orchestrate")
async def orchestrate_endpoint(ticket: TicketRequest):
    """
    Main AI Agent endpoint - fully autonomous:
    - Decides its own workflow
    - Handles multi-turn conversations
    - Auto-escalates when needed
    - Maintains context
    """
    try:
        result = process_with_agent(
            user_message=ticket.text,
            conversation_id=ticket.conversation_id
        )
        
        return {
            "conversation_id": result["conversation_id"],
            "response": result["response"],
            "status": result["status"],
            "message_count": result["message_count"],
            "reasoning_trace": result.get("reasoning_trace", []),
            "instructions": {
                "continue_conversation": "Include the conversation_id in your next request",
                "new_ticket": "Omit conversation_id to start fresh"
            }
        }
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Agent failed: {str(e)}")

# -------------------------------
# Get Conversation History
# -------------------------------
@app.get("/conversation/{conversation_id}")
async def get_conversation(conversation_id: str):
    """Retrieve full conversation history."""
    conv = get_conversation_history(conversation_id)
    if not conv:
        raise HTTPException(status_code=404, detail="Conversation not found")
    
    return {
        "conversation_id": conversation_id,
        "messages": conv["messages"],
        "created_at": conv["created_at"],
        "message_count": len(conv["messages"])
    }

# -------------------------------
# Individual Tool Endpoints (for testing)
# -------------------------------
@app.post("/classify")
async def classify_endpoint(ticket: TicketRequest):
    """Test classification only."""
    classification = classify_ticket(ticket.text)
    return {"classification": classification}

@app.post("/route")
async def route_endpoint(ticket: TicketRequest):
    """Test routing only."""
    department = call_routing(ticket.text)
    return {"department": department}

@app.post("/kb_query")
async def kb_query_endpoint(ticket: TicketRequest):
    """Test KB query only."""
    collection = get_kb_collection()
    if not collection or collection.count() == 0:
        raise HTTPException(status_code=400, detail="KB not set up. Call /setup first.")

    try:
        query_embedding = encoder.encode([ticket.text])[0].tolist()
        result = collection.query(
            query_embeddings=[query_embedding],
            n_results=1,
            include=["documents", "distances", "metadatas"]
        )

        if not result or not result.get('documents') or len(result['documents'][0]) == 0:
            return {"answer": "No relevant KB found.", "confidence": 0.0}

        best_doc = result['documents'][0][0]
        best_distance = result['distances'][0][0] if result.get('distances') else 1.0
        confidence = max(0.0, 1.0 - (best_distance / 2.0))

        return {"answer": best_doc, "confidence": round(float(confidence), 3)}

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"KB query failed: {str(e)}")

# -------------------------------
# Health Check
# -------------------------------
@app.get("/health")
async def health():
    collection = get_kb_collection()
    kb_status = "initialized" if collection and collection.count() > 0 else "not initialized"
    kb_count = collection.count() if collection else 0
    
    return {
        "status": "ok",
        "kb_status": kb_status,
        "kb_records": kb_count,
        "active_conversations": len(conversations),
        "agent_type": "Autonomous ReAct Agent with Gemini"
    }