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
import requests
import torch
import time
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents.react.agent import create_react_agent
from langchain.agents.agent import AgentExecutor

from langchain.tools import Tool
from langchain.prompts import PromptTemplate
import threading
from datetime import datetime
import firebase_admin
from firebase_admin import credentials, firestore
from typing import Optional, Dict, Any

# Environment
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
ROUTING_URL = os.environ.get("ROUTING_URL")
SPACE_URL = os.environ.get("SPACE_URL", "http://localhost:7860")
FIREBASE_CREDS_PATH = os.environ.get("FIREBASE_CREDS_PATH")
firebase_creds_json = os.getenv("FIREBASE_CREDS_JSON")

# Initialize Firebase
db = None
if firebase_creds_json:
    try:
        creds_dict = json.loads(firebase_creds_json)
        cred = credentials.Certificate(creds_dict)

        if not firebase_admin._apps:
            firebase_admin.initialize_app(cred)
        db = firestore.client()
        print("✅ Firebase initialized from FIREBASE_CREDS_JSON")
    except Exception as e:
        import traceback
        print(f"⚠️ Firebase init failed: {e}")
        traceback.print_exc()
else:
    print("⚠️ FIREBASE_CREDS_JSON not found in environment variables")

# Label Dictionary
LABEL_DICTIONARY = {
    "I1": "Low Impact", "I2": "Medium Impact", "I3": "High Impact", "I4": "Critical Impact",
    "U1": "Low Urgency", "U2": "Medium Urgency", "U3": "High Urgency", "U4": "Critical Urgency",
    "T1": "Information", "T2": "Incident", "T3": "Problem", "T4": "Request", "T5": "Question"
}

# Classification Model
clf_model_name = "DavinciTech/BERT_Categorizer"
clf_tokenizer = AutoTokenizer.from_pretrained(clf_model_name, cache_dir="/tmp/transformers")
clf_model = AutoModelForSequenceClassification.from_pretrained(clf_model_name, cache_dir="/tmp/transformers")

def classify_ticket(text):
    """Classify ticket into Impact, Urgency, and Type."""
    inputs = clf_tokenizer(text, return_tensors="pt", truncation=True)
    outputs = clf_model(**inputs)
    logits = outputs.logits[0]
    impact_idx = torch.argmax(logits[:4]).item() + 1
    urgency_idx = torch.argmax(logits[4:8]).item() + 1
    type_idx = torch.argmax(logits[8:]).item() + 1
    return {
        "impact": LABEL_DICTIONARY[f"I{impact_idx}"],
        "urgency": LABEL_DICTIONARY[f"U{urgency_idx}"],
        "type": LABEL_DICTIONARY[f"T{type_idx}"]
    }

# Routing Function
def call_routing(text, retries=3, delay=5):
    """Route ticket to appropriate department."""
    url = ROUTING_URL if ROUTING_URL else f"{SPACE_URL}/route"
    for attempt in range(retries):
        try:
            resp = requests.post(url, json={"text": text}, timeout=30)
            resp.raise_for_status()
            return resp.json().get("department", "General IT")
        except Exception as e:
            print(f"Routing attempt {attempt+1} failed: {e}")
            if attempt < retries - 1:
                time.sleep(delay)
    return "General IT"

# Knowledge Base
CHROMA_PATH = "/tmp/chroma"
COLLECTION_NAME = "knowledge_base"
kb_collection = None
kb_lock = threading.Lock()
encoder = SentenceTransformer("all-MiniLM-L6-v2", cache_folder="/tmp/sentence_transformers")

def get_kb_collection():
    global kb_collection
    if kb_collection is None:
        with kb_lock:
            if kb_collection is None:
                try:
                    chroma_client = chromadb.PersistentClient(
                        path=CHROMA_PATH,
                        settings=Settings(anonymized_telemetry=False, allow_reset=True)
                    )
                    kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
                except Exception as e:
                    print(f"Could not get KB collection: {e}")
    return kb_collection

def query_kb(text: str, top_k: int = 1):
    """Query KB and return answer with confidence."""
    collection = get_kb_collection()
    if not collection or collection.count() == 0:
        return {"answer": None, "confidence": 0.0}
    
    try:
        query_embedding = encoder.encode([text])[0].tolist()
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "distances", "metadatas"]
        )
        
        if not results or not results.get("documents") or len(results["documents"][0]) == 0:
            return {"answer": None, "confidence": 0.0}
        
        answer = results["documents"][0][0]
        distance = results["distances"][0][0] if results.get("distances") else 1.0
        confidence = max(0.0, 1.0 - (distance / 2.0))
        
        return {"answer": answer, "confidence": round(float(confidence), 3)}
    except Exception as e:
        print(f"KB query failed: {e}")
        return {"answer": None, "confidence": 0.0}

# Firestore Helper
def save_ticket_to_firestore(ticket_data: Dict[str, Any]):
    """Save resolved/escalated ticket to Firestore."""
    if not db:
        print("⚠️ Firestore not initialized, skipping save")
        return None
    
    try:
        ticket_ref = db.collection('tickets').document()
        ticket_data['created_at'] = firestore.SERVER_TIMESTAMP
        ticket_data['updated_at'] = firestore.SERVER_TIMESTAMP
        ticket_ref.set(ticket_data)
        print(f"✅ Ticket saved to Firestore: {ticket_ref.id}")
        return ticket_ref.id
    except Exception as e:
        print(f"❌ Firestore save failed: {e}")
        return None

# Gemini LLM
llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=0.3,
    google_api_key=GEMINI_API_KEY
)

# Global conversation storage
conversations = {}

# Tool Functions for Agent
def classify_tool(query: str) -> str:
    """Analyzes ticket severity, impact, urgency, and type. Use when you need to understand ticket priority."""
    result = classify_ticket(query)
    return f"Impact: {result['impact']}, Urgency: {result['urgency']}, Type: {result['type']}"

def routing_tool(query: str) -> str:
    """Identifies which IT department should handle this issue. Use when you need to know responsible team."""
    dept = call_routing(query)
    return f"Department: {dept}"

def kb_tool(query: str) -> str:
    """Searches knowledge base for solutions. Returns answer with confidence score. Use when you need technical solutions."""
    result = query_kb(query)
    if result["answer"] and result["confidence"] > 0.5:
        return f"[KB Confidence: {result['confidence']}]\n{result['answer']}"
    return f"[KB Confidence: {result['confidence']}] No relevant solution found in knowledge base."

def escalation_tool(reason: str) -> str:
    """Creates escalation ticket for human agent. Use ONLY when KB confidence is below 0.6 AND issue is truly complex. Always try KB first!"""
    ticket_id = f"TKT-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    return f"ESCALATED: Ticket {ticket_id} created. Reason: {reason}. Human agent will respond in 2-4 hours."

# Define Tools
tools = [
    Tool(
        name="ClassifyTicket",
        func=classify_tool,
        description="Analyzes ticket to determine impact level, urgency, and type. Use this when you need to understand the severity or priority of an issue."
    ),
    Tool(
        name="RouteTicket",
        func=routing_tool,
        description="Determines which IT department should handle this ticket. Use this when you need to identify the responsible team."
    ),
    Tool(
        name="SearchKnowledgeBase",
        func=kb_tool,
        description="Searches internal knowledge base for solutions. Returns answer with confidence score (0-1). ALWAYS USE THIS FIRST before escalating. Use this when you need to find technical solutions or troubleshooting steps."
    ),
    Tool(
        name="EscalateToHuman",
        func=escalation_tool,
        description="Creates an escalation ticket for human agent review. CRITICAL: Use this ONLY as a LAST RESORT when: 1) KB confidence score is below 0.6 AND you've already tried KB, 2) Issue is extremely complex and unusual, 3) User explicitly confirms the KB solution failed after trying it. DO NOT escalate if KB has a reasonable solution (confidence > 0.6)."
    )
]

# IMPROVED Agent Prompt
AGENT_PROMPT = """You are an intelligent IT Helpdesk AI Agent. Your PRIMARY goal is to resolve tickets using the Knowledge Base. Escalation is a LAST RESORT.

AVAILABLE TOOLS:
{tools}

TOOL NAMES: {tool_names}

CRITICAL RULES:
1. **ALWAYS search Knowledge Base FIRST** - This is your primary tool for resolution
2. **Trust KB solutions with confidence >= 0.6** - These are reliable solutions, provide them to users
3. **ONLY escalate when ABSOLUTELY necessary**:
   - KB confidence is below 0.6 AND no solution found
   - Issue is extremely unusual or complex beyond KB scope
   - User explicitly tried your KB solution and reports it failed
4. **Be thorough with KB** - If first search doesn't work, try rephrasing the query
5. **Maintain context** - Remember conversation history for follow-ups

DECISION WORKFLOW:
NEW TICKET → Search KB → If confidence >= 0.6 → Provide solution → Mark RESOLVED
         ↓
   If confidence < 0.6 → Try rephrasing search → Still low? → Classify & Route → THEN escalate

FOLLOW-UP → Check if user tried solution → Worked? → Mark RESOLVED
                                      ↓
                              Failed? → Search KB again with different query → Still failing? → THEN escalate

IMPORTANT: 
- Don't mention tool names or confidence scores to users
- Provide clear, step-by-step instructions from KB
- Be conversational and helpful
- Escalation means you couldn't solve it - avoid this outcome!

FORMAT:
Question: the user's input
Thought: your reasoning about what to do next
Action: the tool to use (must be one of [{tool_names}])
Action Input: the input for that tool
Observation: the tool's output
... (repeat Thought/Action/Observation as needed)
Thought: I now have enough information to respond
Final Answer: your complete response to the user

Begin!

Question: {input}
Thought: {agent_scratchpad}"""

prompt = PromptTemplate.from_template(AGENT_PROMPT)

# Create Agent
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=6,
    handle_parsing_errors="Check your output and make sure it conforms to the format instructions!",
    return_intermediate_steps=True,
    early_stopping_method="force"
)

# Main Processing Function
def process_with_agent(
    user_message: str, 
    conversation_id: str = None,
    user_email: str = None,
    callback=None
):
    """Process user message through autonomous AI agent."""
    
    if not conversation_id:
        conversation_id = f"conv_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{hash(user_message) % 10000}"
    
    if conversation_id not in conversations:
        conversations[conversation_id] = {
            "messages": [],
            "ticket_info": {},
            "created_at": datetime.now().isoformat(),
            "user_email": user_email,
            "status": "open"
        }
    
    conv = conversations[conversation_id]
    
    conv["messages"].append({
        "role": "user",
        "content": user_message,
        "timestamp": datetime.now().isoformat()
    })
    
    if callback:
        callback({"type": "status", "message": "Agent is thinking..."})
    
    # Build context for follow-ups
    if len(conv["messages"]) > 1:
        context = f"CONVERSATION HISTORY:\n"
        for msg in conv["messages"][-6:-1]:
            context += f"{msg['role'].upper()}: {msg['content']}\n"
        context += f"\nCURRENT MESSAGE: {user_message}"
        agent_input = context
    else:
        agent_input = user_message
    
    try:
        result = agent_executor.invoke({"input": agent_input})
        
        agent_response = result.get("output", "I apologize, I encountered an error.")
        intermediate_steps = result.get("intermediate_steps", [])
        
        # IMPROVED: Check tool usage to determine status
        status = "in_progress"
        should_save = False
        escalated = False
        used_escalation_tool = False
        
        # First, check if EscalateToHuman tool was actually used
        for action, observation in intermediate_steps:
            if action.tool == "EscalateToHuman":
                used_escalation_tool = True
                print("🔍 Detected EscalateToHuman tool usage")
                break
        
        # Determine final status
        if used_escalation_tool:
            # Tool was used - definitely escalated
            status = "escalated"
            should_save = True
            escalated = True
            print("🔴 Ticket ESCALATED (tool used) - Saving to Firebase")
        elif "ESCALATED" in agent_response or "TKT-" in agent_response:
            # Fallback: Check response text
            status = "escalated"
            should_save = True
            escalated = True
            print("🔴 Ticket ESCALATED (text match) - Saving to Firebase")
        elif any(phrase in agent_response.lower() for phrase in [
            "resolved", "you're all set", "should work now", "problem solved", 
            "this should fix", "try these steps", "follow these steps"
        ]):
            # Resolution detected
            status = "resolved"
            should_save = True
            print("✅ Ticket RESOLVED - Saving to Firebase")
        else:
            # Still in progress
            print("⏳ Ticket IN PROGRESS - Not saving yet")
        
        # Extract ticket info from tools
        ticket_info = conv.get("ticket_info", {})
        kb_confidence = None
        escalation_reason = None
        
        for action, observation in intermediate_steps:
            if action.tool == "ClassifyTicket":
                parts = str(observation).split(", ")
                for part in parts:
                    if "Impact:" in part:
                        ticket_info["impact"] = part.split(": ")[1]
                    elif "Urgency:" in part:
                        ticket_info["urgency"] = part.split(": ")[1]
                    elif "Type:" in part:
                        ticket_info["type"] = part.split(": ")[1]
            
            elif action.tool == "RouteTicket":
                ticket_info["department"] = str(observation).replace("Department: ", "")
            
            elif action.tool == "SearchKnowledgeBase":
                # Extract confidence from KB response
                if "[KB Confidence:" in str(observation):
                    try:
                        conf_str = str(observation).split("[KB Confidence: ")[1].split("]")[0]
                        kb_confidence = float(conf_str)
                        ticket_info["kb_confidence"] = kb_confidence
                    except:
                        pass
            
            elif action.tool == "EscalateToHuman":
                # Capture escalation reason
                escalation_reason = action.tool_input
        
        conv["ticket_info"] = ticket_info
        conv["status"] = status
        
        # Build reasoning trace for debugging
        reasoning_trace = []
        for action, observation in intermediate_steps:
            reasoning_trace.append({
                "tool": action.tool,
                "input": action.tool_input,
                "output": str(observation)[:200]
            })
            
            if callback:
                callback({
                    "type": "tool_use",
                    "tool": action.tool,
                    "input": action.tool_input
                })
        
        conv["messages"].append({
            "role": "assistant",
            "content": agent_response,
            "timestamp": datetime.now().isoformat(),
            "reasoning": reasoning_trace
        })
        
        # Save to Firestore if resolved OR escalated
        firestore_id = None
        if should_save:
            print(f"💾 Preparing to save ticket - Status: {status}, Escalated: {escalated}")
            
            firestore_data = {
                "conversation_id": conversation_id,
                "status": status,
                "user_email": user_email or "anonymous",
                "ticket_info": ticket_info,
                "messages": conv["messages"],
                "resolution": agent_response,
                "created_at_iso": conv["created_at"],
                "escalated": escalated,
                "reasoning_trace": reasoning_trace
            }
            
            # Add escalation reason if escalated
            if escalated:
                if escalation_reason:
                    firestore_data["escalation_reason"] = escalation_reason
                elif kb_confidence is not None and kb_confidence < 0.6:
                    firestore_data["escalation_reason"] = f"Low KB confidence: {kb_confidence}"
                else:
                    firestore_data["escalation_reason"] = "Complex issue requiring human intervention"
                
                print(f"🔼 Saving escalated ticket - Reason: {firestore_data['escalation_reason']}")
            
            firestore_id = save_ticket_to_firestore(firestore_data)
            
            if firestore_id:
                print(f"✅ Successfully saved to Firestore with ID: {firestore_id}")
            else:
                print("❌ Failed to save to Firestore")
            
            if callback:
                callback({
                    "type": "saved",
                    "firestore_id": firestore_id,
                    "status": status
                })
        else:
            print(f"⏭️  Not saving - Status: {status}, Should save: {should_save}")
        
        return {
            "conversation_id": conversation_id,
            "response": agent_response,
            "status": status,
            "message_count": len(conv["messages"]),
            "reasoning_trace": reasoning_trace,
            "ticket_info": ticket_info,
            "firestore_id": firestore_id,
            "escalated": escalated
        }
        
    except Exception as e:
        print(f"❌ Agent error: {e}")
        import traceback
        traceback.print_exc()
        
        error_response = "I apologize, I encountered an error. Please try again or I can escalate this to a human agent."
        conv["messages"].append({
            "role": "assistant",
            "content": error_response,
            "timestamp": datetime.now().isoformat()
        })
        
        if callback:
            callback({"type": "error", "message": str(e)})
        
        return {
            "conversation_id": conversation_id,
            "response": error_response,
            "status": "error",
            "error": str(e)
        }

def get_conversation_history(conversation_id: str):
    """Get conversation history."""
    return conversations.get(conversation_id)