import datetime
import os
import re
from typing import TypedDict, Optional, List

from llama_index.core.schema import Document
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from llama_index.core.retrievers import BaseRetriever
from langgraph.graph import StateGraph, END

class GraphState(TypedDict):
    query: str
    response: Optional[str]
    retrieved_docs: Optional[List[Document]]
    context: Optional[str]
    history: list
    retriever: Optional[BaseRetriever]
    llm: any

def _select_chunks_by_relevance_and_position(docs: List[Document], query: str, target_count: int) -> List[Document]:
    """Select chunks based on semantic relevance + positional importance without summarization."""
    if len(docs) <= target_count:
        return docs
    
    selected = []
    if len(docs) >= 2:
        selected = [docs[0], docs[-1]]  
        remaining_slots = target_count - 2
        middle_docs = docs[1:-1]
    else:
        remaining_slots = target_count
        middle_docs = docs[:]
    
    if remaining_slots > 0 and middle_docs:
        query_words = set(query.lower().split())
        scored_docs = []
        
        for doc in middle_docs:
            content = doc.text.lower()
            relevance_score = sum(content.count(word) for word in query_words)
            length_bonus = len(doc.text) / 1000  
            total_score = relevance_score + length_bonus
            scored_docs.append((total_score, doc))
        
        scored_docs.sort(key=lambda x: x[0], reverse=True)
        selected.extend([doc for _, doc in scored_docs[:remaining_slots]])
    
    return selected[:target_count]

def _create_context_with_priorities(docs: List[Document], max_chars: int) -> str:
    """Create context by prioritizing important chunks without summarization."""
    if not docs:
        return ""
    
    contexts = []
    total_chars = 0
    
    priority_docs = []
    
    if len(docs) >= 2:
        priority_docs.append(("BOUNDARY", docs[0]))
        priority_docs.append(("BOUNDARY", docs[-1]))
        
        middle_docs = docs[1:-1] if len(docs) > 2 else []
        middle_docs.sort(key=lambda d: len(d.text), reverse=True)
        priority_docs.extend([("CONTENT", doc) for doc in middle_docs])
    else:
        priority_docs = [("CONTENT", doc) for doc in docs]
    
    for priority_type, doc in priority_docs:
        content = doc.text
        
        if total_chars + len(content) > max_chars:
            remaining_chars = max_chars - total_chars
            if remaining_chars > 200:  
                truncated = content[:remaining_chars]
                last_period = truncated.rfind('.')
                if last_period > remaining_chars * 0.8:  
                    truncated = truncated[:last_period + 1]
                truncated += "...[truncated]"
                contexts.append(f"[{priority_type}] {truncated}")
            break
        
        contexts.append(f"[{priority_type}] {content}")
        total_chars += len(content)
    
    return "\n\n---\n\n".join(contexts)


def router_node(state: GraphState) -> GraphState:
    """
    Router that determines the next step based on available retriever.
    """
    print("---NODE: ROUTER---")
    return state

def general_chat_node(state: GraphState) -> GraphState:
    """
    Generates a response for general conversation using the LLM.
    """
    print("---NODE: GENERAL CHAT---")
    llm = state["llm"]
    query = state["query"]
    history = state["history"]
    history_context = "\n".join([f"Human: {q}\nAI: {a}" for q, a in history])
    current_time = datetime.datetime.now().strftime("%Y-%m-%d %I:%M %p")
    prompt = f"""You are Sarathi, a friendly and knowledgeable AI assistant. 

Instructions:
- The current time is {current_time}.
- Be conversational, helpful, and concise.
- If you don't know the answer to a question, admit it honestly.
- Do not make up information.

**Conversation History:**
{history_context}

**Current User Query:**
Human: {query}
AI:"""

    try:
        response_obj = llm.invoke(prompt)
        response_text = getattr(response_obj, "content", str(response_obj))
    except Exception as e:
        response_text = f"Model inference failed: {str(e)}"

    return {"response": response_text.strip()}

def retrieve_node(state: GraphState) -> GraphState:
    """
    Retrieves relevant documents and creates intelligent context without summarization.
    """
    print("---NODE: RETRIEVE---")
    query = state["query"]
    retriever = state["retriever"]
    history = state["history"]

    context = ""
    retrieved_docs = []
    
    try:
        q_len = len(query.split())
        if q_len < 5:
            top_k = 5  
        elif q_len < 15:
            top_k = 8
        else:
            top_k = 12

        retrieved_docs = retriever.retrieve(query)
        print(f"Retrieved {len(retrieved_docs)} documents")

        if retrieved_docs:
            max_chunks = min(len(retrieved_docs), top_k)
            selected_docs = _select_chunks_by_relevance_and_position(
                retrieved_docs, query, max_chunks
            )
            
            print(f"Selected {len(selected_docs)} chunks for context")
            
            MAX_CONTEXT_CHARS = 6000  
            doc_context = _create_context_with_priorities(selected_docs, MAX_CONTEXT_CHARS)
            
            if history and len(doc_context) < MAX_CONTEXT_CHARS * 0.8:
                history_context = "\n\n".join([f"Human: {q}\nAI: {a}" for q, a in history[-3:]])  # Last 3 exchanges
                remaining_chars = MAX_CONTEXT_CHARS - len(doc_context)
                if len(history_context) <= remaining_chars:
                    context = f"{doc_context}\n\n--- Recent Chat History ---\n{history_context}"
                else:
                    context = doc_context
            else:
                context = doc_context
        else:
            context = "No relevant content found in the document."
            
    except Exception as e:
        print(f"Error in retrieve_node: {e}")
        context = f"Retriever failed: {str(e)}"
        retrieved_docs = []
    
    return {"retrieved_docs": retrieved_docs, "context": context}

def generate_rag_node(state: GraphState) -> GraphState:
    """
    Generates an answer using the retrieved context from the PDF.
    """
    print("---NODE: GENERATE RAG RESPONSE---")
    llm = state["llm"]
    query = state["query"]
    context = state.get("context", "")

    if not context.strip():
        return {"response": "I don't have enough context from the document to answer this question."}

    prompt = f"""You are Sarathi, an intelligent document analysis assistant. Use the provided context to answer the user's question thoroughly and accurately.

Context from document:
{context}

Question: {query}

Instructions:
- Provide detailed, well-structured answers
- If the question has multiple parts, address each part separately
- Cite specific information from the context when possible
- If the context doesn't fully cover the question, clearly state what's missing
- Use a conversational but professional tone

Answer:"""

    try:
        response_obj = llm.invoke(prompt)
        response_text = getattr(response_obj, "content", str(response_obj))
        
    except Exception as e:
        response_text = f"Model inference failed: {str(e)}"

    return {"response": response_text.strip()}


def route_query(state: GraphState) -> str:
    """
    Checks if a retriever is available in the state to decide the next step.
    """
    if state.get("retriever"):
        print("---ROUTING: PDF available, routing to retrieve_node---")
        return "retrieve"
    else:
        print("---ROUTING: No PDF, routing to general_chat_node---")
        return "general_chat"


def build_graph(model_type: str = "groq", retriever=None):
    """
    Builds the workflow graph with LLM and retriever.
    """
    
    def make_llm(temp: float):
        if model_type == "groq":
            api_key = os.getenv("OPENROUTER_API_KEY", "").strip()
            if not api_key:
                raise ValueError("OPENROUTER_API_KEY environment variable not set.")
            return ChatOpenAI(
                model="x-ai/grok-4-fast", 
                base_url="https://openrouter.ai/api/v1", 
                api_key=api_key,
                temperature=temp,
            )
            
        elif model_type == "gemini":
            api_key = os.getenv("GEMINI_API_KEY", "").strip()
            if not api_key:
                raise ValueError("GEMINI_API_KEY environment variable not set.")
            return ChatGoogleGenerativeAI(
                model="gemini-2.0-flash",
                api_key=api_key,
                temperature=temp,
            )
        else:
            raise ValueError("Invalid model_type. Choose 'groq' or 'gemini'.")
    
    workflow = StateGraph(GraphState)

    workflow.add_node("router", router_node)
    workflow.add_node("general_chat", general_chat_node)
    workflow.add_node("retrieve", retrieve_node)
    workflow.add_node("generate", generate_rag_node)

    workflow.set_entry_point("router")

    workflow.add_conditional_edges(
        "router",
        route_query,
        {
            "retrieve": "retrieve",
            "general_chat": "general_chat",
        },
    )

    workflow.add_edge("retrieve", "generate")
    workflow.add_edge("generate", END)
    workflow.add_edge("general_chat", END)

    compiled_graph = workflow.compile()
    
    def graph_wrapper(query: str, temperature: float = 0.7, raw_text: str = None, history=None, retriever_override=None):
        active_retriever = retriever_override or retriever
        llm = make_llm(temperature)
        return compiled_graph.invoke({
            "query": query,
            "retriever": active_retriever.as_retriever() if active_retriever else None,
            "history": history or [],
            "llm": llm,
            "response": None,
            "retrieved_docs": None,
            "context": None,
        })
    
    return graph_wrapper