import os
import time
import gradio as gr
import uvicorn
from fastapi import FastAPI, HTTPException, Depends, File, UploadFile
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from typing import Optional, Dict, Any
import threading
import logging
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.callbacks.base import BaseCallbackHandler
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
import tiktoken

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Configuration ---
CHUNK_SIZE = 800
CHUNK_OVERLAP = 100
MAX_TOKENS = 512
TEMPERATURE = 0.5
RETRIEVAL_K = 5

# --- Token Counting Setup ---
try:
    tokenizer = tiktoken.get_encoding("cl100k_base")
except:
    print("Tiktoken encoder 'cl100k_base' not found. Using basic split().")
    tokenizer = type('obj', (object,), {'encode': lambda x: x.split()})()

def estimate_tokens(text):
    """Estimates token count for a given text."""
    return len(tokenizer.encode(text))

# Custom Callback Handler to track LLM token usage
class TokenUsageCallbackHandler(BaseCallbackHandler):
    """Callback handler to track token usage in LLM calls."""
    def __init__(self):
        super().__init__()
        self.reset_counters()

    def reset_counters(self):
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        self.total_llm_calls = 0

    def on_llm_end(self, response, **kwargs):
        """Collect token usage from the LLM response."""
        self.total_llm_calls += 1
        llm_output = response.llm_output
        
        if llm_output and 'usage_metadata' in llm_output:
            usage = llm_output['usage_metadata']
            prompt_tokens = usage.get('prompt_token_count', 0)
            completion_tokens = usage.get('candidates_token_count', 0)
            
            self.total_prompt_tokens += prompt_tokens
            self.total_completion_tokens += completion_tokens

    def get_total_tokens(self):
        """Returns the total prompt and completion tokens."""
        return {
            "total_prompt_tokens": self.total_prompt_tokens,
            "total_completion_tokens": self.total_completion_tokens,
            "total_llm_tokens": self.total_prompt_tokens + self.total_completion_tokens,
            "total_llm_calls": self.total_llm_calls
        }

# --- Pydantic Models for API ---
class InitializeRequest(BaseModel):
    api_key: str
    document_content: Optional[str] = None

class QueryRequest(BaseModel):
    query: str
    api_key: str

class InitializeResponse(BaseModel):
    success: bool
    message: str
    chunks: Optional[int] = None
    estimated_tokens: Optional[int] = None

class QueryResponse(BaseModel):
    success: bool
    answer: str
    response_time: float
    query_tokens: int
    llm_tokens: Dict[str, int]
    session_stats: Dict[str, int]

class StatsResponse(BaseModel):
    total_queries: int
    total_embedding_tokens: int
    total_llm_tokens: int
    total_llm_calls: int
    initialization_complete: bool

# --- Global Variables ---
class RAGSystem:
    def __init__(self):
        self.vector_store = None
        self.qa_chain = None
        self.token_callback_handler = TokenUsageCallbackHandler()
        self.session_stats = {
            "total_queries": 0,
            "total_embedding_tokens": 0,
            "initialization_complete": False
        }
        self.current_api_key = None

# Global RAG system instance
rag_system = RAGSystem()

def initialize_rag_system(api_key, file_content=None):
    """Initialize the RAG system with API key and optional file content."""
    global rag_system
    
    try:
        # Set API key
        os.environ["GOOGLE_API_KEY"] = api_key
        rag_system.current_api_key = api_key
        
        # Initialize embeddings
        embeddings = GoogleGenerativeAIEmbeddings(
            model="models/embedding-001",
            google_api_key=api_key
        )
        
        # Initialize LLM
        llm = ChatGoogleGenerativeAI(
            model="gemini-1.5-flash",
            google_api_key=api_key,
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            callbacks=[rag_system.token_callback_handler],
            verbose=False
        )
        
        # Load or use default document
        if file_content:
            # Save uploaded file content
            with open("uploaded_document.txt", "w", encoding="utf-8") as f:
                f.write(file_content)
            loader = TextLoader("uploaded_document.txt")
        else:
            # Check if default maize_data.txt exists
            if os.path.exists("maize_data.txt"):
                loader = TextLoader("maize_data.txt")
            else:
                return "❌ No document found. Please upload a file or ensure maize_data.txt exists."
        
        # Load and split documents
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=CHUNK_SIZE, 
            chunk_overlap=CHUNK_OVERLAP
        )
        chunks = text_splitter.split_documents(documents)
        
        # Estimate embedding tokens
        initial_embedding_tokens = sum(estimate_tokens(chunk.page_content) for chunk in chunks)
        rag_system.session_stats["total_embedding_tokens"] = initial_embedding_tokens
        
        # Create vector store
        rag_system.vector_store = FAISS.from_documents(chunks, embeddings)
        
        # Create prompt template
        prompt_template = PromptTemplate(
            input_variables=["context", "question"],
            template="""
You are an expert in maize agriculture. Use the following context ONLY to answer the question accurately and helpfully. If the context doesn't contain the answer, say "Based on the provided context, I cannot answer this question.".

Context:
{context}

Question: {question}

Answer:"""
        )
        
        # Set up QA chain
        rag_system.qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=rag_system.vector_store.as_retriever(search_kwargs={"k": RETRIEVAL_K}),
            chain_type_kwargs={"prompt": prompt_template},
            callbacks=[rag_system.token_callback_handler],
            return_source_documents=True
        )
        
        rag_system.session_stats["initialization_complete"] = True
        
        return f"✅ RAG system initialized successfully!\n📄 Document processed: {len(chunks)} chunks\n🔢 Estimated embedding tokens: ~{initial_embedding_tokens}"
        
    except Exception as e:
        logger.error(f"Initialization failed: {str(e)}")
        return f"❌ Initialization failed: {str(e)}"

def process_query(query, api_key):
    """Process a user query through the RAG system."""
    global rag_system
    
    if not api_key:
        return "❌ Please provide a Google API key first.", ""
    
    if not rag_system.qa_chain:
        return "❌ RAG system not initialized. Please initialize first.", ""
    
    if not query.strip():
        return "❌ Please enter a question.", ""
    
    try:
        # Estimate query embedding tokens
        query_tokens = estimate_tokens(query)
        rag_system.session_stats["total_embedding_tokens"] += query_tokens
        rag_system.session_stats["total_queries"] += 1
        
        # Process query
        start_time = time.time()
        result = rag_system.qa_chain({"query": query})
        end_time = time.time()
        
        # Get token usage
        llm_tokens = rag_system.token_callback_handler.get_total_tokens()
        
        # Format response
        answer = result['result']
        
        # Create stats summary
        stats = f"""
📊 **Query Statistics:**
- Response time: {end_time - start_time:.2f} seconds
- Query tokens (estimated): ~{query_tokens}
- LLM tokens (this query): Prompt: {llm_tokens['total_prompt_tokens']}, Completion: {llm_tokens['total_completion_tokens']}

📈 **Session Statistics:**
- Total queries: {rag_system.session_stats['total_queries']}
- Total embedding tokens: ~{rag_system.session_stats['total_embedding_tokens']}
- Total LLM calls: {llm_tokens['total_llm_calls']}
- Total LLM tokens: {llm_tokens['total_llm_tokens']}
"""
        
        return answer, stats
        
    except Exception as e:
        logger.error(f"Error processing query: {str(e)}")
        return f"❌ Error processing query: {str(e)}", ""

def upload_file_and_initialize(api_key, file):
    """Handle file upload and system initialization."""
    if not api_key:
        return "❌ Please provide a Google API key first."
    
    if file is None:
        return initialize_rag_system(api_key)
    
    try:
        # Handle different file object types based on Gradio version
        if hasattr(file, 'name'):
            # Newer Gradio versions - file has .name attribute
            with open(file.name, 'r', encoding='utf-8') as f:
                file_content = f.read()
        elif isinstance(file, str):
            # File path as string
            with open(file, 'r', encoding='utf-8') as f:
                file_content = f.read()
        elif hasattr(file, 'read'):
            # File-like object
            file_content = file.read()
            if isinstance(file_content, bytes):
                file_content = file_content.decode('utf-8')
        else:
            # Fallback - try to read as bytes and decode
            file_content = file.decode('utf-8') if isinstance(file, bytes) else str(file)
            
        return initialize_rag_system(api_key, file_content)
        
    except Exception as e:
        logger.error(f"Error reading uploaded file: {str(e)}")
        return f"❌ Error reading uploaded file: {str(e)}"

def reset_session():
    """Reset the session statistics."""
    global rag_system
    rag_system.token_callback_handler.reset_counters()
    rag_system.session_stats = {
        "total_queries": 0,
        "total_embedding_tokens": 0,
        "initialization_complete": False
    }
    return "🔄 Session statistics reset."

# --- FastAPI Setup ---
app = FastAPI(
    title="Maize RAG Q&A System API",
    description="API for the Maize Agriculture RAG Q&A System",
    version="1.0.0"
)

# Optional: Add API key authentication for API endpoints
security = HTTPBearer(auto_error=False)

async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """Extract API key from Authorization header (optional)"""
    if credentials:
        return credentials.credentials
    return None

# --- API Endpoints ---

@app.get("/")
async def root():
    """Root endpoint"""
    return {"message": "Maize RAG Q&A System API", "status": "running"}

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "system_initialized": rag_system.session_stats["initialization_complete"]
    }

@app.post("/initialize", response_model=InitializeResponse)
async def initialize_system(request: InitializeRequest):
    """Initialize the RAG system"""
    try:
        result = initialize_rag_system(request.api_key, request.document_content)
        
        if "✅" in result:
            # Parse successful result
            lines = result.split('\n')
            chunks = None
            tokens = None
            
            for line in lines:
                if "chunks" in line:
                    chunks = int(line.split(': ')[1].split(' ')[0])
                elif "tokens" in line:
                    tokens = int(line.split('~')[1])
            
            return InitializeResponse(
                success=True,
                message=result,
                chunks=chunks,
                estimated_tokens=tokens
            )
        else:
            return InitializeResponse(
                success=False,
                message=result
            )
    
    except Exception as e:
        logger.error(f"API initialization error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/query", response_model=QueryResponse)
async def query_system(request: QueryRequest):
    """Query the RAG system"""
    try:
        if not rag_system.session_stats["initialization_complete"]:
            raise HTTPException(status_code=400, detail="System not initialized")
        
        # Estimate query embedding tokens
        query_tokens = estimate_tokens(request.query)
        rag_system.session_stats["total_embedding_tokens"] += query_tokens
        rag_system.session_stats["total_queries"] += 1
        
        # Process query
        start_time = time.time()
        result = rag_system.qa_chain({"query": request.query})
        end_time = time.time()
        
        # Get token usage
        llm_tokens = rag_system.token_callback_handler.get_total_tokens()
        
        response_time = end_time - start_time
        
        return QueryResponse(
            success=True,
            answer=result['result'],
            response_time=response_time,
            query_tokens=query_tokens,
            llm_tokens=llm_tokens,
            session_stats=rag_system.session_stats
        )
    
    except Exception as e:
        logger.error(f"API query error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/stats", response_model=StatsResponse)
async def get_stats():
    """Get current session statistics"""
    llm_tokens = rag_system.token_callback_handler.get_total_tokens()
    
    return StatsResponse(
        total_queries=rag_system.session_stats["total_queries"],
        total_embedding_tokens=rag_system.session_stats["total_embedding_tokens"],
        total_llm_tokens=llm_tokens["total_llm_tokens"],
        total_llm_calls=llm_tokens["total_llm_calls"],
        initialization_complete=rag_system.session_stats["initialization_complete"]
    )

@app.post("/reset")
async def reset_system():
    """Reset session statistics"""
    reset_session()
    return {"message": "Session reset successfully"}

@app.post("/upload-document")
async def upload_document(
    file: UploadFile = File(...),
    api_key: str = None
):
    """Upload a document and initialize the system"""
    try:
        if not api_key:
            raise HTTPException(status_code=400, detail="API key required")
        
        # Read uploaded file
        content = await file.read()
        file_content = content.decode('utf-8')
        
        # Initialize system with uploaded content
        result = initialize_rag_system(api_key, file_content)
        
        if "✅" in result:
            return {"success": True, "message": result}
        else:
            return {"success": False, "message": result}
    
    except Exception as e:
        logger.error(f"Document upload error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# Create Gradio interface with version compatibility
def create_interface():
    # Check Gradio version for compatibility
    import gradio as gr
    gradio_version = gr.__version__
    
    with gr.Blocks(title="Maize RAG Q&A System", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🌽 Maize Agriculture RAG Q&A System
        
        This system uses Retrieval-Augmented Generation (RAG) to answer questions about maize agriculture.
        Upload your own document or use the default maize dataset.
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                api_key_input = gr.Textbox(
                    label="🔑 Google API Key",
                    placeholder="Enter your Google Generative AI API key",
                    type="password"
                )
                gr.Markdown("Get your API key from Google AI Studio")
            
            with gr.Column(scale=1):
                reset_btn = gr.Button("🔄 Reset Session", variant="secondary")
        
        with gr.Row():
            with gr.Column():
                file_upload = gr.File(
                    label="📁 Upload Document (Optional)",
                    file_types=[".txt"]
                )
                gr.Markdown("Upload a text file or use the default maize dataset")
                
                init_btn = gr.Button("🚀 Initialize RAG System", variant="primary")
                init_output = gr.Textbox(
                    label="📋 Initialization Status",
                    lines=3,
                    interactive=False
                )
        
        gr.Markdown("## 💬 Ask Questions")
        
        with gr.Row():
            with gr.Column(scale=3):
                query_input = gr.Textbox(
                    label="❓ Your Question",
                    placeholder="Ask something about maize agriculture...",
                    lines=2
                )
                
                # Sample questions
                sample_questions = [
                    "What are the main pests affecting maize crops?",
                    "How should maize be irrigated?",
                    "What is the ideal soil type for maize?",
                    "What are the nutritional requirements of maize?",
                    "When is the best time to harvest maize?"
                ]
                
                # Use Examples component if available, otherwise just show as markdown
                try:
                    gr.Examples(
                        examples=sample_questions,
                        inputs=query_input,
                        label="💡 Sample Questions"
                    )
                except:
                    gr.Markdown("💡 **Sample Questions:**\n" + 
                               "\n".join([f"- {q}" for q in sample_questions]))
            
            with gr.Column(scale=1):
                submit_btn = gr.Button("🔍 Ask", variant="primary")
        
        with gr.Row():
            with gr.Column(scale=2):
                answer_output = gr.Textbox(
                    label="🤖 Answer",
                    lines=6,
                    interactive=False
                )
            
            with gr.Column(scale=1):
                stats_output = gr.Markdown(
                    value="📊 Statistics will appear here after queries."
                )
        
        # Event handlers
        init_btn.click(
            upload_file_and_initialize,
            inputs=[api_key_input, file_upload],
            outputs=init_output
        )
        
        submit_btn.click(
            process_query,
            inputs=[query_input, api_key_input],
            outputs=[answer_output, stats_output]
        )
        
        query_input.submit(
            process_query,
            inputs=[query_input, api_key_input],
            outputs=[answer_output, stats_output]
        )
        
        reset_btn.click(
            reset_session,
            outputs=init_output
        )
        
        gr.Markdown("""
        ## 📝 Instructions:
        1. **Enter your Google API Key** (required)
        2. **Upload a document** (optional - uses default maize dataset if not provided)
        3. **Initialize the RAG system** by clicking "Initialize RAG System"
        4. **Ask questions** about the document content
        5. **View statistics** to monitor token usage and costs
        
        ## 💰 Cost Information:
        - **Gemini 1.5 Flash**: Input: $0.075/1M tokens, Output: $0.30/1M tokens
        - **Embedding Model**: $0.025/1M tokens
        
        Token usage is estimated and displayed for cost tracking.
        
        ## 🔗 API Access:
        This system also provides REST API endpoints:
        - **API Docs**: Add `/docs` to the URL for interactive API documentation
        - **Health Check**: `GET /health`
        - **Initialize**: `POST /initialize`
        - **Query**: `POST /query`
        """)
    
    return demo

# Create and launch the interface
def run_gradio():
    """Run Gradio interface"""
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        quiet=True  # Reduce Gradio logs in combined mode
    )

def run_fastapi():
    """Run FastAPI server"""
    uvicorn.run(
        app, 
        host="0.0.0.0", 
        port=8000,
        log_level="info"
    )

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        mode = sys.argv[1]
        
        if mode == "api":
            # Run only FastAPI
            print("Starting FastAPI server on port 8000...")
            run_fastapi()
        elif mode == "gradio":
            # Run only Gradio
            print("Starting Gradio interface on port 7860...")
            run_gradio()
        elif mode == "both":
            # Run both servers
            print("Starting both FastAPI (port 8000) and Gradio (port 7860)...")
            
            # Start FastAPI in a separate thread
            fastapi_thread = threading.Thread(target=run_fastapi)
            fastapi_thread.daemon = True
            fastapi_thread.start()
            
            # Start Gradio in main thread
            time.sleep(2)  # Give FastAPI time to start
            run_gradio()
        else:
            print("Usage: python app.py [api|gradio|both]")
            print("Default: gradio only")
            run_gradio()
    else:
        # Default: run only Gradio (for Hugging Face Spaces compatibility)
        print("Starting Gradio interface on port 7860...")
        run_gradio()