"""
Deep Research Agent - AI-Powered Research Assistant
Conducts comprehensive web research and delivers detailed reports via email.
"""

import gradio as gr
import asyncio
import os
import re
import json
import traceback
from typing import List, Dict, Tuple, Optional

from langchain_openai import ChatOpenAI
from langchain_community.tools import DuckDuckGoSearchResults
from pydantic import BaseModel, Field
import sendgrid
from sendgrid.helpers.mail import Mail, Email, To, Content


# ===========================
# Data Models
# ===========================

class WebSearchItem(BaseModel):
    """Represents a single web search query with reasoning."""
    reason: str = Field(description="Reasoning for why this search is important")
    query: str = Field(description="The search term to use")


class WebSearchPlan(BaseModel):
    """Collection of planned web searches."""
    searches: List[WebSearchItem] = Field(description="List of searches to perform")


class ReportData(BaseModel):
    """Structure for the final research report."""
    short_summary: str = Field(description="2-3 sentence summary")
    markdown_report: str = Field(description="Full markdown report")
    follow_up_questions: List[str] = Field(description="Suggested research topics")


# ===========================
# Configuration & Validation
# ===========================

class Config:
    """Application configuration."""
    NUM_SEARCHES = 3
    SUMMARY_MAX_WORDS = 300
    REPORT_MIN_WORDS = 1000
    MODEL_NAME = "gpt-4o-mini"
    MODEL_TEMPERATURE = 0.7


def validate_email(email: str) -> bool:
    """Validate email format using regex."""
    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
    return bool(re.match(pattern, email))


def validate_openai_key(api_key: str) -> bool:
    """Validate OpenAI API key format."""
    return api_key.startswith('sk-') and len(api_key) > 20


def check_sendgrid_config() -> Tuple[bool, str]:
    """Check if SendGrid is properly configured."""
    sendgrid_key = os.environ.get('SENDGRID_API_KEY')
    from_email = os.environ.get('FROM_EMAIL')
    
    if not sendgrid_key or not from_email:
        return False, "⚠️ SendGrid not configured. Contact administrator."
    return True, ""


# ===========================
# Research Planning
# ===========================

def plan_searches(query: str, llm: ChatOpenAI) -> WebSearchPlan:
    """Generate a plan of web searches for the given query."""
    prompt = f"""You are a research assistant. Plan {Config.NUM_SEARCHES} web searches for this query.

Query: {query}

Respond with valid JSON in this format:
{{
    "searches": [
        {{"reason": "reasoning", "query": "search term"}},
        {{"reason": "reasoning", "query": "search term"}},
        {{"reason": "reasoning", "query": "search term"}}
    ]
}}

Only output JSON, nothing else."""

    try:
        response = llm.invoke(prompt)
        content = extract_json(response.content)
        return WebSearchPlan(**json.loads(content))
    except Exception:
        # Fallback to default searches
        return WebSearchPlan(searches=[
            WebSearchItem(reason="Primary research", query=query),
            WebSearchItem(reason="Recent developments", query=f"{query} latest 2025"),
            WebSearchItem(reason="Expert analysis", query=f"{query} trends analysis")
        ])


def extract_json(content: str) -> str:
    """Extract JSON from markdown code blocks or raw text."""
    if "```json" in content:
        return content.split("```json")[1].split("```")[0].strip()
    elif "```" in content:
        return content.split("```")[1].split("```")[0].strip()
    return content


# ===========================
# Web Search
# ===========================

def perform_web_search(query: str) -> str:
    """Execute a DuckDuckGo web search."""
    try:
        search = DuckDuckGoSearchResults()
        return search.run(query)
    except Exception as e:
        return f"Search failed: {str(e)}"


def summarize_search_results(query: str, reason: str, results: str, llm: ChatOpenAI) -> str:
    """Generate a concise summary of search results."""
    prompt = f"""Summarize these search results concisely.

Search: {query}
Purpose: {reason}
Results: {results}

Create a 2-3 paragraph summary under {Config.SUMMARY_MAX_WORDS} words. Focus on facts and key insights. No commentary."""

    response = llm.invoke(prompt)
    return response.content


async def conduct_research(search_plan: WebSearchPlan, llm: ChatOpenAI, progress) -> List[str]:
    """Execute all searches and generate summaries."""
    summaries = []
    total = len(search_plan.searches)
    
    for i, item in enumerate(search_plan.searches):
        progress((i + 1) / total, desc=f"Searching: {item.query}")
        
        results = perform_web_search(item.query)
        summary = summarize_search_results(item.query, item.reason, results, llm)
        summaries.append(summary)
    
    return summaries


# ===========================
# Report Generation
# ===========================

def generate_report(query: str, summaries: List[str], llm: ChatOpenAI) -> ReportData:
    """Synthesize research into a comprehensive report."""
    summaries_text = "\n\n".join([
        f"Search {i+1} Results:\n{summary}" 
        for i, summary in enumerate(summaries)
    ])
    
    prompt = f"""Create a comprehensive research report.

Query: {query}

Research Findings:
{summaries_text}

Generate a detailed report (minimum {Config.REPORT_MIN_WORDS} words) in markdown format.

Respond with valid JSON:
{{
    "short_summary": "2-3 sentence overview",
    "markdown_report": "# Title\\n\\nFull report...",
    "follow_up_questions": ["question 1", "question 2", "question 3"]
}}

Only output JSON."""

    try:
        response = llm.invoke(prompt)
        content = extract_json(response.content)
        return ReportData(**json.loads(content))
    except Exception:
        # Fallback report
        report_text = f"# Research Report: {query}\n\n" + "\n\n".join(summaries)
        return ReportData(
            short_summary=f"Research findings on {query}",
            markdown_report=report_text,
            follow_up_questions=["Further investigation needed", "Related topics", "Deep dive areas"]
        )


# ===========================
# Email Delivery
# ===========================

def markdown_to_html(markdown: str, llm: ChatOpenAI) -> str:
    """Convert markdown to professional HTML email."""
    prompt = f"""Convert this markdown to beautiful HTML for email.

{markdown}

Create professional HTML with proper styling, headings, and formatting.
Start with <!DOCTYPE html>. Only output HTML."""

    response = llm.invoke(prompt)
    return response.content


def send_email_report(report: ReportData, to_email: str, llm: ChatOpenAI) -> Tuple[bool, str]:
    """Send research report via email."""
    try:
        sendgrid_key = os.environ.get('SENDGRID_API_KEY')
        from_email_addr = os.environ.get('FROM_EMAIL')
        
        # Validate configuration
        if not sendgrid_key:
            return False, "❌ SendGrid API key not configured in environment variables"
        if not from_email_addr:
            return False, "❌ FROM_EMAIL not configured in environment variables"
        
        # Prepare email content
        subject = f"Research Report: {report.short_summary[:50]}..."
        html_body = markdown_to_html(report.markdown_report, llm)
        
        # Initialize SendGrid
        sg = sendgrid.SendGridAPIClient(api_key=sendgrid_key)
        
        # Create mail object
        mail = Mail(
            from_email=Email(from_email_addr),
            to_emails=To(to_email),
            subject=subject,
            html_content=Content("text/html", html_body)
        )
        
        # Send email
        response = sg.client.mail.send.post(request_body=mail.get())
        
        if response.status_code in [200, 201, 202]:
            return True, f"✅ Email sent successfully! (Status: {response.status_code})\n📬 Check your inbox (and spam folder)"
        else:
            return False, f"⚠️ SendGrid returned status {response.status_code}"
            
    except Exception as e:
        error_msg = str(e)
        
        # Provide specific guidance based on error
        if "403" in error_msg or "Forbidden" in error_msg:
            return False, f"""❌ SendGrid Error 403: Access Forbidden

This means your SendGrid configuration needs attention:

🔍 **Check These Issues:**

1. **API Key Permissions**
   - Go to SendGrid → Settings → API Keys
   - Your API key must have "Mail Send" permission
   - Try creating a NEW API key with "Full Access"

2. **Sender Email Verification** (Most Common Issue!)
   - Go to SendGrid → Settings → Sender Authentication
   - Your FROM_EMAIL ({os.environ.get('FROM_EMAIL', 'not set')}) must show as "Verified" ✅
   - If not verified, click "Verify a Single Sender"
   - Check your email and click the verification link
   
3. **API Key Format**
   - Your key should start with 'SG.'
   - Copy the FULL key when setting up

4. **Environment Variables in Hugging Face**
   - Go to Space Settings → Repository Secrets
   - Make sure both secrets are saved:
     * SENDGRID_API_KEY = SG.xxxxx...
     * FROM_EMAIL = your-verified-email@domain.com

**Next Steps:**
1. Verify your sender email in SendGrid
2. Create a new API key with full permissions
3. Update both secrets in Hugging Face
4. Restart your Space

Original error: {error_msg}"""
        
        elif "401" in error_msg or "Unauthorized" in error_msg:
            return False, f"""❌ SendGrid Error 401: Invalid API Key

Your API key is not valid.

**Fix:**
1. Go to SendGrid → Settings → API Keys
2. Create a NEW API key
3. Copy the full key (starts with SG.)
4. Update SENDGRID_API_KEY in Hugging Face Space settings
5. Restart your Space

Original error: {error_msg}"""
        
        else:
            return False, f"❌ Email sending failed: {error_msg}\n\n{traceback.format_exc()}"


# ===========================
# Main Research Pipeline
# ===========================

async def execute_research(query: str, llm: ChatOpenAI, progress) -> Tuple[Optional[ReportData], Optional[str]]:
    """Execute the complete research pipeline."""
    try:
        # Plan
        progress(0.1, desc="Planning research strategy...")
        search_plan = plan_searches(query, llm)
        
        # Research
        progress(0.3, desc="Conducting web searches...")
        summaries = await conduct_research(search_plan, llm, progress)
        
        # Synthesize
        progress(0.8, desc="Generating comprehensive report...")
        report = generate_report(query, summaries, llm)
        
        progress(1.0, desc="Complete!")
        return report, None
        
    except Exception as e:
        return None, f"Research failed: {str(e)}"


def process_research_request(openai_key: str, user_email: str, query: str, progress=gr.Progress()) -> Tuple[str, str, str, str]:
    """Main entry point for research requests."""
    
    # Validate inputs
    if not openai_key or not openai_key.strip():
        return "❌ OpenAI API key required", "", "", "❌ Missing API key"
    
    if not validate_openai_key(openai_key):
        return "❌ Invalid OpenAI API key format", "", "", "❌ Invalid API key"
    
    if not user_email or not user_email.strip():
        return "❌ Email address required", "", "", "❌ Missing email"
    
    if not validate_email(user_email):
        return "❌ Invalid email format", "", "", "❌ Invalid email"
    
    if not query or not query.strip():
        return "❌ Research query required", "", "", "❌ Missing query"
    
    # Check SendGrid
    sendgrid_configured, config_error = check_sendgrid_config()
    
    try:
        # Initialize LLM
        os.environ['OPENAI_API_KEY'] = openai_key
        llm = ChatOpenAI(model=Config.MODEL_NAME, temperature=Config.MODEL_TEMPERATURE)
        
        # Execute research
        progress(0.0, desc="Initializing research...")
        report, error = asyncio.run(execute_research(query, llm, progress))
        
        if error:
            return error, "", "", f"❌ {error}"
        
        # Prepare outputs
        markdown_report = report.markdown_report
        summary = report.short_summary
        followup = "\n".join([f"- {q}" for q in report.follow_up_questions])
        
        # Send email
        email_status = ""
        if sendgrid_configured:
            try:
                progress(0.95, desc=f"Sending to {user_email}...")
                success, message = send_email_report(report, user_email, llm)
                email_status = message if success else f"⚠️ {message}"
            except Exception as e:
                email_status = f"⚠️ Email error: {str(e)}"
        else:
            email_status = config_error or "✅ Research complete (Email not configured)"
        
        return markdown_report, summary, followup, email_status
        
    except Exception as e:
        return f"Error: {str(e)}", "", "", f"❌ {str(e)}"


# ===========================
# UI Components
# ===========================

CUSTOM_CSS = """
/* Modal Overlay */
.modal-overlay {
    position: fixed !important;
    top: 0 !important;
    left: 0 !important;
    right: 0 !important;
    bottom: 0 !important;
    background: rgba(0, 0, 0, 0.75) !important;
    display: flex !important;
    justify-content: center !important;
    align-items: center !important;
    z-index: 9999 !important;
    backdrop-filter: blur(4px);
}

/* Modal Box */
.modal-box {
    background: white !important;
    border-radius: 16px !important;
    padding: 2.5rem !important;
    box-shadow: 0 25px 50px -12px rgba(0, 0, 0, 0.5) !important;
    max-width: 500px !important;
    width: 90% !important;
    animation: modalSlideIn 0.3s ease-out !important;
}

@keyframes modalSlideIn {
    from {
        opacity: 0;
        transform: translateY(-20px) scale(0.95);
    }
    to {
        opacity: 1;
        transform: translateY(0) scale(1);
    }
}

/* Modal Header */
.modal-header {
    text-align: center !important;
    margin-bottom: 1.5rem !important;
}

.modal-icon {
    font-size: 3rem !important;
    margin-bottom: 1rem !important;
}

.modal-title {
    font-size: 1.5rem !important;
    font-weight: 700 !important;
    color: #1a1a1a !important;
    margin-bottom: 0.5rem !important;
}

.modal-subtitle {
    font-size: 0.95rem !important;
    color: #666 !important;
}

.step-indicator {
    display: inline-block !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    padding: 0.25rem 0.75rem !important;
    border-radius: 20px !important;
    font-size: 0.75rem !important;
    font-weight: 600 !important;
    margin-bottom: 1rem !important;
}

/* Buttons */
.primary-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    padding: 0.75rem 2rem !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    cursor: pointer !important;
    transition: all 0.2s !important;
    width: 100% !important;
}

.primary-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4) !important;
}

.secondary-btn {
    background: #f3f4f6 !important;
    border: 1px solid #d1d5db !important;
    color: #374151 !important;
    padding: 0.75rem 1.5rem !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    cursor: pointer !important;
    transition: all 0.2s !important;
}

.secondary-btn:hover {
    background: #e5e7eb !important;
}

/* Hide main content when modal is shown */
.hidden {
    display: none !important;
}

/* Main interface styling */
.main-interface {
    max-width: 1200px;
    margin: 0 auto;
    padding: 2rem;
}

.settings-btn {
    float: right;
    background: #f3f4f6;
    border: 1px solid #d1d5db;
    padding: 0.5rem 1rem;
    border-radius: 6px;
    cursor: pointer;
    transition: all 0.2s;
}

.settings-btn:hover {
    background: #e5e7eb;
}
"""


# ===========================
# Gradio Interface
# ===========================

def create_interface():
    """Create the Gradio interface with popup modals."""
    
    with gr.Blocks(title="Deep Research Agent", css=CUSTOM_CSS, theme=gr.themes.Soft()) as app:
        
        # State management
        user_email_state = gr.State("")
        openai_key_state = gr.State("")
        
        # EMAIL MODAL
        with gr.Group(elem_classes="modal-overlay", visible=True) as email_modal:
            with gr.Group(elem_classes="modal-box"):
                with gr.Column():
                    gr.HTML('<div class="modal-header"><div class="modal-icon">📧</div></div>')
                    gr.HTML('<div class="step-indicator">STEP 1 OF 2</div>')
                    gr.Markdown('<h2 class="modal-title">Welcome to Deep Research Agent</h2>')
                    gr.Markdown('<p class="modal-subtitle">Enter your email to receive research reports</p>')
                    
                    email_input = gr.Textbox(
                        label="Email Address",
                        placeholder="your.email@example.com",
                        lines=1,
                        show_label=False
                    )
                    email_error = gr.Markdown("", visible=False, elem_classes="error-message")
                    
                    email_continue_btn = gr.Button(
                        "Continue →",
                        elem_classes="primary-btn",
                        size="lg"
                    )
        
        # API KEY MODAL
        with gr.Group(elem_classes="modal-overlay", visible=False) as api_modal:
            with gr.Group(elem_classes="modal-box"):
                with gr.Column():
                    gr.HTML('<div class="modal-header"><div class="modal-icon">🔑</div></div>')
                    gr.HTML('<div class="step-indicator">STEP 2 OF 2</div>')
                    gr.Markdown('<h2 class="modal-title">Enter Your OpenAI API Key</h2>')
                    gr.Markdown('<p class="modal-subtitle">Get your key at <a href="https://platform.openai.com" target="_blank">platform.openai.com</a></p>')
                    
                    api_input = gr.Textbox(
                        label="OpenAI API Key",
                        placeholder="sk-...",
                        type="password",
                        lines=1,
                        show_label=False
                    )
                    api_error = gr.Markdown("", visible=False, elem_classes="error-message")
                    
                    with gr.Row():
                        api_back_btn = gr.Button(
                            "← Back",
                            elem_classes="secondary-btn",
                            scale=1
                        )
                        api_continue_btn = gr.Button(
                            "Start Researching →",
                            elem_classes="primary-btn",
                            scale=2
                        )
        
        # MAIN INTERFACE
        with gr.Group(visible=False, elem_classes="main-interface") as main_interface:
            # Header with settings
            with gr.Row():
                gr.Markdown("# 🔍 Deep Research Agent")
                change_settings_btn = gr.Button("⚙️ Change Settings", elem_classes="settings-btn", size="sm")
            
            gr.Markdown("### AI-Powered Research Assistant")
            gr.Markdown("*Ask any research question and receive a comprehensive report via email*")
            
            query_input = gr.Textbox(
                label="What would you like to research?",
                placeholder="E.g., Latest developments in quantum computing",
                lines=3
            )
            
            research_btn = gr.Button("🚀 Start Research", variant="primary", size="lg")
            
            status_output = gr.Textbox(
                label="📮 Status",
                interactive=False,
                lines=5,
                max_lines=10
            )
            
            with gr.Accordion("💡 Example Research Topics", open=False):
                gr.Markdown("""
                - Latest developments in quantum computing
                - Impact of AI on healthcare in 2025
                - Sustainable energy solutions comparison
                - Recent breakthroughs in gene therapy
                - Future of autonomous vehicles technology
                - Comparison of large language models 2025
                - Advances in renewable energy storage
                - Emerging trends in biotechnology
                """)
            
            gr.Markdown("---")
            
            with gr.Row():
                with gr.Column():
                    summary_output = gr.Textbox(
                        label="📝 Executive Summary",
                        lines=4,
                        show_copy_button=True
                    )
                with gr.Column():
                    followup_output = gr.Textbox(
                        label="🔗 Follow-up Questions",
                        lines=4,
                        show_copy_button=True
                    )
            
            gr.Markdown("### 📄 Full Research Report")
            report_output = gr.Markdown()
            
            # Footer
            gr.Markdown("---")
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("""
                    **💰 Cost**  
                    ~$0.05-0.10 per query
                    """)
                with gr.Column(scale=1):
                    gr.Markdown("""
                    **⏱️ Duration**  
                    1-3 minutes
                    """)
                with gr.Column(scale=1):
                    gr.Markdown("""
                    **🔒 Privacy**  
                    Never stored
                    """)
                with gr.Column(scale=1):
                    gr.Markdown("""
                    **🔍 Search**  
                    Free DuckDuckGo
                    """)
        
        # ===========================
        # Event Handlers
        # ===========================
        
        def validate_email_step(email):
            """Validate email and move to API key modal."""
            if not email or not email.strip():
                return {
                    email_error: gr.update(value="❌ Please enter your email address", visible=True),
                }
            if not validate_email(email):
                return {
                    email_error: gr.update(value="❌ Please enter a valid email address", visible=True),
                }
            return {
                user_email_state: email,
                email_error: gr.update(visible=False),
                email_modal: gr.update(visible=False),
                api_modal: gr.update(visible=True)
            }
        
        def validate_api_step(api_key, email):
            """Validate API key and show main interface."""
            if not api_key or not api_key.strip():
                return {
                    api_error: gr.update(value="❌ Please enter your OpenAI API key", visible=True),
                }
            if not validate_openai_key(api_key):
                return {
                    api_error: gr.update(value="❌ Invalid API key format. Should start with 'sk-'", visible=True),
                }
            return {
                openai_key_state: api_key,
                api_error: gr.update(visible=False),
                api_modal: gr.update(visible=False),
                main_interface: gr.update(visible=True)
            }
        
        def go_back_to_email():
            """Return to email modal."""
            return {
                api_modal: gr.update(visible=False),
                email_modal: gr.update(visible=True)
            }
        
        def show_settings():
            """Show settings (email modal) from main interface."""
            return {
                main_interface: gr.update(visible=False),
                email_modal: gr.update(visible=True)
            }
        
        # Connect events
        email_continue_btn.click(
            fn=validate_email_step,
            inputs=[email_input],
            outputs=[user_email_state, email_error, email_modal, api_modal]
        )
        
        api_continue_btn.click(
            fn=validate_api_step,
            inputs=[api_input, user_email_state],
            outputs=[openai_key_state, api_error, api_modal, main_interface]
        )
        
        api_back_btn.click(
            fn=go_back_to_email,
            outputs=[api_modal, email_modal]
        )
        
        change_settings_btn.click(
            fn=show_settings,
            outputs=[main_interface, email_modal]
        )
        
        research_btn.click(
            fn=process_research_request,
            inputs=[openai_key_state, user_email_state, query_input],
            outputs=[report_output, summary_output, followup_output, status_output]
        )
    
    return app


# ===========================
# Application Entry Point
# ===========================

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()