# ===================================================================
# AI Research Agent - Agentic RAG System for Hugging Face Spaces
# ===================================================================
import os
import re
import json
import ast
import operator
import logging
import requests
import tempfile
import time
import asyncio
from pathlib import Path
from typing import List, Dict, Any, Optional
from datetime import datetime
from urllib.parse import quote_plus

# Core Libraries
import numpy as np
import pandas as pd
from tqdm import tqdm

# ML & Embedding
import PyPDF2
from sentence_transformers import SentenceTransformer
import faiss

# LLM & Web
import groq
from groq import Groq

# UI & Voice
import gradio as gr
from gtts import gTTS
try:
    import speech_recognition as sr
    STT_AVAILABLE = True
except ImportError:
    STT_AVAILABLE = False

GTTS_AVAILABLE = True

# ===================================================================
# CONFIGURATION & LOGGING
# ===================================================================
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ===================================================================
# UTILITY CLASSES
# ===================================================================

class WebSearchTool:
    def __init__(self, max_results: int = 5, timeout: int = 10):
        self.max_results = max_results
        self.timeout = timeout
        self.base_url = "https://api.duckduckgo.com/"

    def search(self, query: str, num_results: Optional[int] = None) -> Dict[str, Any]:
        num_results = num_results or self.max_results
        try:
            params = {
                'q': query,
                'format': 'json',
                'no_redirect': '1',
                'no_html': '1',
                'skip_disambig': '1'
            }
            response = requests.get(self.base_url, params=params, timeout=self.timeout,
                                  headers={'User-Agent': 'AI Research Agent 1.0'})
            response.raise_for_status()
            data = response.json()

            results = {
                'query': query,
                'abstract': data.get('Abstract', ''),
                'abstract_source': data.get('AbstractSource', ''),
                'answer': data.get('Answer', ''),
                'related_topics': [],
                'results_found': bool(any([data.get('Abstract'), data.get('Answer')]))
            }

            if 'RelatedTopics' in data:
                for topic in data['RelatedTopics'][:num_results]:
                    if isinstance(topic, dict) and 'Text' in topic:
                        results['related_topics'].append({
                            'text': topic.get('Text', ''),
                            'url': topic.get('FirstURL', '')
                        })
            return results
        except Exception as e:
            logger.error(f"Web search failed: {e}")
            return {'query': query, 'error': str(e), 'results_found': False}


class ConfigManager:
    DEFAULT_CONFIG = {
        'embedding_model': 'all-MiniLM-L6-v2',
        'groq_model': 'llama-3.1-8b-instant',
        'max_iterations': 5,
        'confidence_threshold': 0.7,
        'retrieval_k': 5,
        'chunk_size': 512,
        'chunk_overlap': 50
    }

    @staticmethod
    def load_config():
        return ConfigManager.DEFAULT_CONFIG.copy()


# ===================================================================
# DOCUMENT PROCESSING
# ===================================================================

class DocumentProcessor:
    def __init__(self):
        self.supported_extensions = {'.txt', '.md', '.pdf'}

    def load_documents(self, data_directory: str) -> List[Dict[str, Any]]:
        documents = []
        data_path = Path(data_directory)
        if not data_path.exists():
            return documents

        files = [f for f in data_path.rglob('*') if f.suffix.lower() in self.supported_extensions]
        for file_path in tqdm(files, desc="Loading documents"):
            try:
                content = self._extract_text(file_path)
                if content.strip():
                    doc = {
                        'doc_id': str(file_path.relative_to(data_path)),
                        'content': content,
                        'file_path': str(file_path),
                        'file_type': file_path.suffix.lower()
                    }
                    documents.append(doc)
            except Exception as e:
                logger.error(f"Error loading {file_path}: {e}")
        return documents

    def _extract_text(self, file_path: Path) -> str:
        extension = file_path.suffix.lower()
        if extension == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        elif extension == '.pdf':
            text = ""
            with open(file_path, 'rb') as f:
                pdf_reader = PyPDF2.PdfReader(f)
                for page in pdf_reader.pages:
                    text += page.extract_text() + "\n"
            return text
        return ""


class DocumentChunker:
    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def chunk_documents(self, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        chunks = []
        for doc in tqdm(documents, desc="Chunking documents"):
            doc_chunks = self._split_text(doc['content'])
            for i, chunk_text in enumerate(doc_chunks):
                chunk = {
                    'chunk_id': f"{doc['doc_id']}_chunk_{i}",
                    'content': chunk_text,
                    'doc_id': doc['doc_id'],
                    'chunk_index': i,
                    'source_file': doc['file_path'],
                    'file_type': doc['file_type']
                }
                chunks.append(chunk)
        return chunks

    def _split_text(self, text: str) -> List[str]:
        text = re.sub(r'\s+', ' ', text.strip())
        if len(text) <= self.chunk_size:
            return [text]

        chunks = []
        start = 0
        while start < len(text):
            end = start + self.chunk_size
            if end >= len(text):
                chunks.append(text[start:])
                break

            chunk = text[start:end]
            last_sentence = max(chunk.rfind('.'), chunk.rfind('!'), chunk.rfind('?'))
            if last_sentence > start + self.chunk_size // 2:
                end = start + last_sentence + 1
            else:
                last_space = chunk.rfind(' ')
                if last_space > start + self.chunk_size // 2:
                    end = start + last_space

            chunks.append(text[start:end].strip())
            start = end - self.chunk_overlap

        return [chunk for chunk in chunks if len(chunk.strip()) > 10]


class EmbeddingGenerator:
    def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
        self.model_name = model_name
        self.model = SentenceTransformer(model_name)

    def generate_embeddings(self, chunks: List[Dict[str, Any]]) -> np.ndarray:
        texts = [chunk['content'] for chunk in chunks]
        embeddings = self.model.encode(texts, batch_size=32, show_progress_bar=True, convert_to_numpy=True)
        return embeddings

    def get_query_embedding(self, query: str) -> np.ndarray:
        return self.model.encode([query], convert_to_numpy=True)[0]


def build_embeddings_from_directory(data_directory: str, output_directory: str,
                                   chunk_size: int = 512, chunk_overlap: int = 50) -> Dict[str, Any]:
    os.makedirs(output_directory, exist_ok=True)
    doc_processor = DocumentProcessor()
    chunker = DocumentChunker(chunk_size, chunk_overlap)
    embedder = EmbeddingGenerator()

    documents = doc_processor.load_documents(data_directory)
    if not documents:
        return {}

    chunks = chunker.chunk_documents(documents)
    embeddings = embedder.generate_embeddings(chunks)

    return {
        'chunks': chunks,
        'embeddings': embeddings,
        'metadata': {
            'num_documents': len(documents),
            'num_chunks': len(chunks),
            'embedding_dim': embeddings.shape[1]
        }
    }


# ===================================================================
# RETRIEVER
# ===================================================================

class DocumentRetriever:
    def __init__(self, embedding_model_name: str = 'all-MiniLM-L6-v2'):
        self.embedding_generator = EmbeddingGenerator(embedding_model_name)
        self.index = None
        self.chunks = []
        self.embeddings = None

    def build_index(self, chunks: List[Dict[str, Any]], embeddings: np.ndarray) -> None:
        self.chunks = chunks
        self.embeddings = embeddings
        embedding_dim = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(embedding_dim)
        embeddings_normalized = self._normalize_embeddings(embeddings)
        self.index.add(embeddings_normalized.astype(np.float32))

    def _normalize_embeddings(self, embeddings: np.ndarray) -> np.ndarray:
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms[norms == 0] = 1
        return embeddings / norms

    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        if not self.index:
            return []

        query_embedding = self.embedding_generator.get_query_embedding(query)
        query_embedding_normalized = self._normalize_embeddings(query_embedding.reshape(1, -1))
        scores, indices = self.index.search(query_embedding_normalized.astype(np.float32), k)

        results = []
        for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
            if idx >= 0:
                chunk = self.chunks[idx].copy()
                chunk.update({'similarity_score': float(score), 'rank': i + 1})
                results.append(chunk)
        return results


# ===================================================================
# AGENTIC TOOLS
# ===================================================================

class AgenticTools:
    def __init__(self):
        self.tools = {
            "calculator": self.calculator_tool,
            "web_search": self.web_search_tool,
            "fact_checker": self.fact_checker_tool,
            "document_analyzer": self.document_analyzer_tool
        }
        self.web_search_instance = WebSearchTool()

    def calculator_tool(self, expression: str) -> Dict[str, Any]:
        try:
            clean_expr = re.sub(r'[^0-9+\-*/().\s]', '', expression)
            node = ast.parse(clean_expr, mode='eval')
            result = self._eval_expr(node.body)
            return {
                "tool": "calculator",
                "input": expression,
                "result": result,
                "success": True,
                "explanation": f"Calculated {clean_expr} = {result}"
            }
        except Exception as e:
            return {"tool": "calculator", "input": expression, "result": None, "success": False, "error": str(e)}

    def _eval_expr(self, node):
        ops = {
            ast.Add: operator.add, ast.Sub: operator.sub,
            ast.Mult: operator.mul, ast.Div: operator.truediv,
            ast.Pow: operator.pow, ast.USub: operator.neg
        }
        if isinstance(node, ast.Num):
            return node.n
        elif isinstance(node, ast.BinOp):
            return ops[type(node.op)](self._eval_expr(node.left), self._eval_expr(node.right))
        elif isinstance(node, ast.UnaryOp):
            return ops[type(node.op)](self._eval_expr(node.operand))
        raise TypeError(node)

    def web_search_tool(self, query: str) -> Dict[str, Any]:
        try:
            result = self.web_search_instance.search(query)
            return {
                "tool": "web_search",
                "input": query,
                "result": result,
                "success": result.get('results_found', False),
                "explanation": f"Found web information about: {query}"
            }
        except Exception as e:
            return {"tool": "web_search", "input": query, "result": None, "success": False, "error": str(e)}

    def fact_checker_tool(self, claim: str) -> Dict[str, Any]:
        confidence = "medium"
        verification = "partial"
        if re.search(r'\d+', claim):
            verification = "requires_calculation"
        return {
            "tool": "fact_checker",
            "input": claim,
            "result": {"verification": verification, "confidence": confidence},
            "success": True
        }

    def document_analyzer_tool(self, text: str, analysis_type: str = "summary") -> Dict[str, Any]:
        sentences = re.split(r'[.!?]+', text)[:3]
        summary = '. '.join([s.strip() for s in sentences if s.strip()])
        return {
            "tool": "document_analyzer",
            "input": f"{analysis_type} analysis",
            "result": summary,
            "success": True
        }


class AgentPlanner:
    def __init__(self):
        self.planning_patterns = {
            "calculation": ["calculate", "compute", "math", "percentage", "total"],
            "current_info": ["latest", "recent", "current", "rate", "price", "exchange", "dollar", "currency"],
            "analysis": ["analyze", "insights", "patterns", "summary"],
            "fact_check": ["verify", "confirm", "accurate"]
        }

    def create_execution_plan(self, query: str) -> Dict[str, Any]:
        query_lower = query.lower()
        needed_capabilities = []
        for capability, keywords in self.planning_patterns.items():
            if any(keyword in query_lower for keyword in keywords):
                needed_capabilities.append(capability)

        steps = [{"step": 1, "tool": "document_search", "description": "Search documents", "query": query}]
        step_num = 2

        if "calculation" in needed_capabilities:
            steps.append({"step": step_num, "tool": "calculator", "description": "Perform calculations", "depends_on": [1]})
            step_num += 1
        if "current_info" in needed_capabilities:
            steps.append({"step": step_num, "tool": "web_search", "description": "Search web", "query": query, "depends_on": [1]})
            step_num += 1
        if "analysis" in needed_capabilities:
            steps.append({"step": step_num, "tool": "document_analyzer", "description": "Analyze content", "depends_on": [1]})
            step_num += 1

        steps.append({"step": step_num, "tool": "synthesizer", "description": "Synthesize results", "depends_on": list(range(1, step_num))})

        return {"query": query, "detected_needs": needed_capabilities, "steps": steps, "total_steps": len(steps)}


class ResultSynthesizer:
    def __init__(self, groq_client):
        self.groq_client = groq_client

    def synthesize_results(self, query: str, results: Dict[str, Any], temperature: float = 0.3, max_tokens: int = 500) -> str:
        context_parts = []
        if "document_search" in results and results["document_search"]["success"]:
            context_parts.append(f"DOCUMENTS:\n{results['document_search']['result']}")
        if "web_search" in results and results["web_search"]["success"]:
            web_info = results["web_search"]["result"]
            web_text = f"{web_info.get('abstract', '')} {web_info.get('answer', '')}"
            context_parts.append(f"WEB INFO:\n{web_text}")
        if "calculator" in results and results["calculator"]["success"]:
            context_parts.append(f"CALCULATION:\n{results['calculator']['result']}")

        all_context = "\n\n".join(context_parts)
        prompt = f"""Based on the following information, provide a comprehensive answer.
QUESTION: {query}
INFORMATION:
{all_context}
Provide a clear, direct answer synthesizing all sources."""

        try:
            response = self.groq_client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {"role": "system", "content": "You are an expert research assistant."},
                    {"role": "user", "content": prompt}
                ],
                temperature=temperature,
                max_tokens=max_tokens
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            return f"Based on available information: {all_context[:500]}..."


class AgenticEvaluator:
    def evaluate_response(self, query: str, response: str, tool_results: Dict[str, Any]) -> Dict[str, Any]:
        successful_tools = sum(1 for r in tool_results.values() if r.get("success", False))
        total_tools = len(tool_results)

        confidence = min(0.8, successful_tools / max(total_tools, 1)) if successful_tools > 0 else 0.0
        source_types = []
        if "document_search" in tool_results and tool_results["document_search"]["success"]:
            source_types.append("documents")
        if "web_search" in tool_results and tool_results["web_search"]["success"]:
            source_types.append("web")

        return {
            "confidence_score": confidence,
            "completeness": "comprehensive" if successful_tools >= total_tools else "partial",
            "source_diversity": len(source_types),
            "recommendations": []
        }


# ===================================================================
# MAIN AGENT CLASS
# ===================================================================

class AgenticRAGAgent:
    def __init__(self):
        self.config = ConfigManager.load_config()
        self.retriever = None
        self.groq_client = None
        self.conversation_history = []

        self.tools = AgenticTools()
        self.planner = AgentPlanner()
        self.synthesizer = None
        self.evaluator = AgenticEvaluator()

        self.temperature = 0.3
        self.max_tokens = 500
        self.chunk_size = 512
        self.chunk_overlap = 50
        self.retrieval_k = 8

        self.enable_web_search = True
        self.enable_calculations = True
        self.enable_fact_checking = True
        self.enable_analysis = True

        # Initialize Groq
        groq_api_key = os.getenv("GROQ_API_KEY")
        if groq_api_key:
            try:
                self.groq_client = Groq(api_key=groq_api_key)
                self.synthesizer = ResultSynthesizer(self.groq_client)
                print("✅ Groq API configured")
            except Exception as e:
                print(f"❌ Error: {e}")

    def clean_text_for_speech(self, text):
        """Clean text for TTS"""
        if not text:
            return ""

        # Remove markdown formatting
        text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text)
        text = re.sub(r'\*([^*]+)\*', r'\1', text)
        text = re.sub(r'^#{1,6}\s+', '', text, flags=re.MULTILINE)
        text = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text)
        text = re.sub(r'```[^`]*```', '', text, flags=re.DOTALL)
        text = re.sub(r'`([^`]+)`', r'\1', text)
        text = re.sub(r'^[\s]*[-*+•]\s+', '', text, flags=re.MULTILINE)
        text = re.sub(r'^[\s]*\d+\.\s+', '', text, flags=re.MULTILINE)

        # Remove emojis
        emoji_pattern = re.compile(
            "["
            "\U0001F600-\U0001F64F"
            "\U0001F300-\U0001F5FF"
            "\U0001F680-\U0001F6FF"
            "\U0001F1E0-\U0001F1FF"
            "\U00002702-\U000027B0"
            "\U000024C2-\U0001F251"
            "\U0001F900-\U0001F9FF"
            "\U00002600-\U000026FF"
            "\U00002700-\U000027BF"
            "]+"
        )
        text = emoji_pattern.sub('', text)
        text = re.sub(r'\s+', ' ', text)
        text = re.sub(r'\n+', '. ', text)
        text = text.strip()
        text = re.sub(r'\.+', '.', text)

        return text

    def generate_audio_response(self, text):
        """Generate audio using gTTS"""
        if not text or not GTTS_AVAILABLE:
            return None

        clean_text = self.clean_text_for_speech(text)
        if not clean_text:
            return None

        try:
            temp_dir = tempfile.gettempdir()
            timestamp = int(time.time())
            audio_file = os.path.join(temp_dir, f"response_{timestamp}.mp3")

            tts = gTTS(text=clean_text, lang='en', slow=False)
            tts.save(audio_file)
            return audio_file
        except Exception as e:
            logger.error(f"Audio generation failed: {e}")
            return None

    def is_greeting_or_casual(self, query):
        query_lower = query.lower().strip()
        greetings = ['hi', 'hello', 'hey', 'howdy']
        return any(query_lower.startswith(g) for g in greetings) or query_lower in greetings

    def get_greeting_response(self, query):
        return "Hi there! 👋 I'm AI Research Agent with agentic capabilities. Upload PDF documents and ask complex questions!"

    def get_simple_answer(self, query, retrieved_docs):
        if not self.groq_client:
            return "Error: Groq API not configured"

        context = "\n\n".join([doc.get('content', str(doc)) for doc in retrieved_docs[:5]])
        prompt = f"""Based on this context, provide a clear answer.
Context: {context}
Question: {query}
Answer:"""

        try:
            response = self.groq_client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {"role": "system", "content": "You are a helpful research assistant."},
                    {"role": "user", "content": prompt}
                ],
                temperature=self.temperature,
                max_tokens=self.max_tokens
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            return f"Error: {str(e)}"

    async def process_agentic_query(self, query, chat_history, progress=gr.Progress()):
        if not query.strip():
            return chat_history, "", None

        if chat_history is None:
            chat_history = []

        chat_history.append({"role": "user", "content": query})

        try:
            if self.is_greeting_or_casual(query):
                progress(0.5, desc="Generating response...")
                response = self.get_greeting_response(query)
                chat_history.append({"role": "assistant", "content": response})

                progress(0.8, desc="🔊 Generating voice...")
                audio_file = self.generate_audio_response(response)

                return chat_history, "", audio_file

            progress(0.1, desc="🧠 Planning...")

            if not self.retriever or not hasattr(self.retriever, 'index') or not self.retriever.index:
                error = "📄 Please upload a PDF document first!"
                chat_history.append({"role": "assistant", "content": error})
                audio_file = self.generate_audio_response(error)
                return chat_history, "", audio_file

            plan = self.planner.create_execution_plan(query)
            progress(0.2, desc=f"📋 Plan: {len(plan['steps'])} steps")

            results = {}
            current_step = 0

            for step in plan['steps']:
                current_step += 1
                progress_val = 0.2 + (current_step / len(plan['steps'])) * 0.6
                progress(progress_val, desc=f"🔧 Step {current_step}: {step['description']}")

                if step['tool'] == 'document_search':
                    retrieved_docs = self.retriever.search(query, k=self.retrieval_k)
                    if retrieved_docs:
                        doc_answer = self.get_simple_answer(query, retrieved_docs)
                        results['document_search'] = {"success": True, "result": doc_answer}
                    else:
                        results['document_search'] = {"success": False, "result": "No relevant info"}

                elif step['tool'] == 'calculator' and self.enable_calculations:
                    math_patterns = re.findall(r'[\d+\-*/().\s]+', query)
                    for expr in math_patterns:
                        if any(op in expr for op in ['+', '-', '*', '/']):
                            results['calculator'] = self.tools.calculator_tool(expr.strip())
                            break

                elif step['tool'] == 'web_search' and self.enable_web_search:
                    results['web_search'] = self.tools.web_search_tool(query)

                elif step['tool'] == 'document_analyzer' and self.enable_analysis:
                    if 'document_search' in results and results['document_search']['success']:
                        doc_content = results['document_search']['result']
                        results['document_analyzer'] = self.tools.document_analyzer_tool(doc_content, "summary")

            progress(0.85, desc="🔬 Synthesizing...")

            if self.synthesizer:
                final_answer = self.synthesizer.synthesize_results(query, results, self.temperature, self.max_tokens)
            else:
                successful = [r['result'] for r in results.values() if r.get('success')]
                final_answer = f"Based on available info: {' '.join(map(str, successful))}"

            progress(0.9, desc="📊 Evaluating...")
            evaluation = self.evaluator.evaluate_response(query, final_answer, results)

            eval_summary = f"\n\n💡 **Analysis:**\n"
            eval_summary += f"• Confidence: {evaluation['confidence_score']:.1%}\n"
            eval_summary += f"• Sources: {evaluation['source_diversity']} types\n"
            eval_summary += f"• Completeness: {evaluation['completeness']}"

            complete_response = final_answer + eval_summary

            progress(0.95, desc="🔊 Generating voice response...")
            audio_file = self.generate_audio_response(final_answer)

            chat_history.append({"role": "assistant", "content": complete_response})

            self.conversation_history.append({
                'timestamp': datetime.now().isoformat(),
                'query': query,
                'response': complete_response,
                'plan': plan,
                'results': results,
                'evaluation': evaluation,
                'audio_file': audio_file
            })

            progress(1.0, desc="✅ Complete!")
            return chat_history, "", audio_file

        except Exception as e:
            error = f"❌ Error: {str(e)}"
            chat_history.append({"role": "assistant", "content": error})
            return chat_history, "", None

    def upload_documents(self, files, progress=gr.Progress()):
        if not files:
            return "No files uploaded"

        try:
            progress(0.1, desc="Processing files...")
            os.makedirs("sample_data", exist_ok=True)

            uploaded = []
            for file in files:
                if hasattr(file, 'name') and file.name.endswith('.pdf'):
                    original = os.path.basename(file.name)
                    dest = os.path.join("sample_data", original)
                    with open(dest, "wb") as dst:
                        dst.write(file.read())

                    uploaded.append(original)

            if not uploaded:
                return "❌ No valid PDF files"

            progress(0.5, desc="Generating embeddings...")
            embeddings_data = build_embeddings_from_directory("sample_data", "temp_embeddings")

            if embeddings_data and 'embeddings' in embeddings_data:
                progress(0.8, desc="Building index...")
                self.retriever = DocumentRetriever()
                self.retriever.build_index(embeddings_data['chunks'], embeddings_data['embeddings'])

                doc_count = embeddings_data.get('metadata', {}).get('num_documents', 0)
                chunk_count = embeddings_data.get('metadata', {}).get('num_chunks', 0)

                progress(1.0, desc="Complete!")
                return f"""✅ **Success!**
📄 Files: {', '.join(uploaded)}
📊 Documents: {doc_count} | Chunks: {chunk_count}
🎯 Ready for complex questions with voice support!"""
            else:
                return "❌ Failed to process documents"
        except Exception as e:
            return f"❌ Error: {str(e)}"

    def update_settings(self, temp, tokens, chunk_size, overlap, k, web, calc, fact, analysis):
        self.temperature = temp
        self.max_tokens = tokens
        self.chunk_size = chunk_size
        self.chunk_overlap = overlap
        self.retrieval_k = k
        self.enable_web_search = web
        self.enable_calculations = calc
        self.enable_fact_checking = fact
        self.enable_analysis = analysis

        return f"""⚙️ Settings Updated:
• Temperature: {temp}
• Max Tokens: {tokens}
• Chunk Size: {chunk_size}
• Retrieved: {k}
• Web: {'✅' if web else '❌'}
• Calc: {'✅' if calc else '❌'}
• Voice Output: {'✅' if GTTS_AVAILABLE else '❌'}"""


# ===================================================================
# GRADIO INTERFACE (COMPATIBLE WITH GRADIO 4.27)
# ===================================================================

def create_interface():
    agent = AgenticRAGAgent()

    with gr.Blocks(title="🤖 AI Research Agent", theme=gr.themes.Soft()) as interface:
        gr.HTML("""
        <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px;">
            <h1 style="color: white; margin: 0;">🤖 AI Research Agent - Agentic RAG</h1>
            <p style="color: white; margin: 10px 0;">Advanced Multi-Tool Research Assistant with Voice Support 🔊</p>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(label="💬 Chat", height=500)

                with gr.Row():
                    msg = gr.Textbox(label="", placeholder="Ask a complex research question...", scale=4)
                    submit_btn = gr.Button("🚀 Send", variant="primary", scale=1)

                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")

                audio_output = gr.Audio(label="🔊 Voice Response", autoplay=True, interactive=False)

            with gr.Column(scale=1):
                with gr.Group():
                    gr.HTML("<h3 style='text-align: center;'>📄 Upload Documents</h3>")
                    file_upload = gr.Files(label="", file_types=[".pdf"], file_count="multiple")
                upload_status = gr.Textbox(label="📊 Status", interactive=False, max_lines=10)

                with gr.Accordion("⚙️ Settings", open=False):
                    gr.HTML("<h4>🧠 AI Parameters</h4>")
                    temperature_slider = gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="🌡️ Temperature")
                    max_tokens_slider = gr.Slider(100, 1000, value=500, step=50, label="📝 Max Tokens")

                    gr.HTML("<h4>📄 Document Processing</h4>")
                    chunk_size_slider = gr.Slider(256, 1024, value=512, step=64, label="📄 Chunk Size")
                    chunk_overlap_slider = gr.Slider(0, 100, value=50, step=10, label="🔗 Overlap")
                    retrieval_k_slider = gr.Slider(3, 15, value=8, step=1, label="🔍 Retrieved Chunks")

                    gr.HTML("<h4>🛠️ Agentic Tools</h4>")
                    with gr.Row():
                        enable_web = gr.Checkbox(value=True, label="🌐 Web Search")
                        enable_calc = gr.Checkbox(value=True, label="🧮 Calculator")
                    with gr.Row():
                        enable_fact = gr.Checkbox(value=True, label="✅ Fact Check")
                        enable_analysis = gr.Checkbox(value=True, label="📊 Analysis")

                    apply_btn = gr.Button("⚡ Apply Settings", variant="primary", size="lg")

                settings_status = gr.Textbox(label="⚙️ Settings Status", interactive=False, max_lines=8)

                with gr.Accordion("🔊 Voice Features Status", open=False):
                    gr.HTML(f"""
                    <div style="padding: 10px;">
                        <p><strong>Text-to-Speech (gTTS):</strong> {'✅ Available' if GTTS_AVAILABLE else '❌ Not Available'}</p>
                        <p><strong>Speech-to-Text:</strong> {'✅ Available' if STT_AVAILABLE else '❌ Not Available (HF Spaces limitation)'}</p>
                        <p><em>Voice output: Auto-plays with responses</em></p>
                    </div>
                    """)

        # -----------------------------
        # Event Handlers (Sync wrapper for async)
        # -----------------------------
        def process_msg(message, history):
            import asyncio
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    future = asyncio.run_coroutine_threadsafe(agent.process_agentic_query(message, history), loop)
                    return future.result()
                else:
                    return loop.run_until_complete(agent.process_agentic_query(message, history))
            except RuntimeError:
                return asyncio.run(agent.process_agentic_query(message, history))

        submit_btn.click(process_msg, inputs=[msg, chatbot], outputs=[chatbot, msg, audio_output])
        msg.submit(process_msg, inputs=[msg, chatbot], outputs=[chatbot, msg, audio_output])
        clear_btn.click(lambda: [], outputs=[chatbot])

        file_upload.change(agent.upload_documents, inputs=[file_upload], outputs=[upload_status])

        apply_btn.click(
            agent.update_settings,
            inputs=[
                temperature_slider, max_tokens_slider, chunk_size_slider,
                chunk_overlap_slider, retrieval_k_slider, enable_web,
                enable_calc, enable_fact, enable_analysis
            ],
            outputs=[settings_status]
        )

    return interface

# ===================================================================
# MAIN
# ===================================================================

if __name__ == "__main__":
    print("🚀 Launching AI Research Agent on Hugging Face Spaces...")
    print("✨ Features:")
    print("   • Multi-Tool Integration")
    print("   • Intelligent Query Planning")
    print("   • Multi-Step Reasoning")
    print("   • Result Synthesis")
    print("   • Quality Evaluation")
    print("   • 🔊 Voice Output (Text-to-Speech)")

    app = create_interface()
    app.launch()