import logging
import json
import asyncio
from typing import Optional, Dict, Any
from langchain_community.llms import Ollama
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from core.config import OPENAI_API_KEY
import chromadb

logger = logging.getLogger("chatbot.ai_service")

# ── Vector DB Initialization ────────────────────────────────────────────────────
try:
    chroma_client = chromadb.PersistentClient(path="./chroma_data")
except Exception as e:
    logger.error(f"ChromaDB initialization failed: {e}")
    chroma_client = None

def get_vector_collection(collection_name: str = "default_knowledge"):
    if not chroma_client: return None
    return chroma_client.get_or_create_collection(name=collection_name)

# ── AI Model Initialization ────────────────────────────────────────────────────
local_llm = None
try:
    local_llm = Ollama(model="llama3")
    logger.info("Local AI (Ollama) initialized.")
except Exception as e:
    logger.warning(f"Local AI init failed: {e}")

cloud_llm = None
if OPENAI_API_KEY:
    try:
        cloud_llm = ChatOpenAI(temperature=0.3, openai_api_key=OPENAI_API_KEY, model="gpt-4o-mini")
        logger.info("Cloud AI (OpenAI) initialized.")
    except Exception as e:
        logger.warning(f"Cloud AI init failed: {e}")

def _build_system_prompt(custom_prompt: Optional[str], context: str) -> str:
    base = custom_prompt or (
        "You are an enterprise AI assistant.\n"
        "Answer the user's question accurately using ONLY the provided [CONTEXT].\n"
        "If you don't know, say so."
    )
    return f"{base}\n\n[CONTEXT]\n{context}"

# ── Core Fallback Generation ───────────────────────────────────────────────────

async def generate_chatbot_response(
    query: str,
    context: str,
    custom_prompt: Optional[str] = None
) -> Dict[str, Any]:
    """
    Implements the core Local AI -> Cloud AI fallback chain.
    """
    system_prompt = _build_system_prompt(custom_prompt, context)
    full_prompt = f"{system_prompt}\n\nUser: {query}\nAssistant:"

    # 1. Attempt Local AI
    if local_llm:
        try:
            logger.info("Attempting Local AI generation...")
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(None, lambda: local_llm.invoke(full_prompt))
            return {"response": response, "provider": "local_ollama", "tokens_used": len(response.split())}
        except Exception as e:
            logger.warning(f"Local AI failed: {e}. Falling back to Cloud.")
    
    # 2. Attempt Cloud AI Fallback
    if cloud_llm:
        try:
            logger.info("Attempting Cloud AI fallback...")
            from langchain.schema import HumanMessage, SystemMessage
            messages = [
                SystemMessage(content=system_prompt),
                HumanMessage(content=query)
            ]
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(None, lambda: cloud_llm.invoke(messages))
            
            tokens = getattr(response.response_metadata.get('token_usage', {}), 'total_tokens', 0) if hasattr(response, 'response_metadata') else 0
            
            return {"response": response.content, "provider": "cloud_openai", "tokens_used": tokens}
        except Exception as e:
            logger.error(f"Cloud AI failed: {e}")

    # 3. Ultimate Fallback (Error State)
    return {
        "response": "The AI systems (both local and cloud) are currently unavailable.",
        "provider": "error",
        "tokens_used": 0
    }

# ── Vector Search Methods ──────────────────────────────────────────────────────

def store_document_vector(document_id: str, text: str, metadata: dict, collection_name: str = "default_knowledge"):
    """Stores text embeddings into ChromaDB for semantic search."""
    try:
        collection = get_vector_collection(collection_name)
        if not collection: return False
        collection.add(
            documents=[text],
            metadatas=[metadata],
            ids=[document_id]
        )
        return True
    except Exception as e:
        logger.error(f"Failed to store vector: {e}")
        return False

def semantic_search(query: str, n_results: int = 3, collection_name: str = "default_knowledge"):
    """Performs a semantic search against ChromaDB."""
    try:
        collection = get_vector_collection(collection_name)
        if not collection: return None
        results = collection.query(
            query_texts=[query],
            n_results=n_results
        )
        return results
    except Exception as e:
        logger.error(f"Vector search failed: {e}")
        return None