import aiohttp
import asyncio
import os
from dotenv import load_dotenv
import logging
from typing import List, Dict
import xml.etree.ElementTree as ET
import ast

from gmail_mcp.utils.GCP.gmail_search import search_emails
from .prompts import email_response_analysis, draft_prompt_modification, email_response_requirements_analysis, draft_response_generation
from ..GCP.gmail_indexer import GmailIndexer
import base64

logger = logging.getLogger("deepsearch-mcp")

# Load environment variables from .env file
load_dotenv()

# Get API key from environment variables
OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY')

# Add these constants at the top with your other imports
API_URL = "https://openrouter.ai/api/v1/chat/completions"


async def make_api_call(prompt, model="openai/o3-mini", temperature=0.1, session=None, max_retries=5, base_delay=1):
    """Make API call with exponential backoff retry logic."""
    logger.info(f"Making API call to OpenRouter with model: {model}")
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "HTTP-Referer": "http://localhost:8000",
    }
    data = {
        "messages": [{"role": "user", "content": prompt}],
        "model": model,
        "max_tokens": 100000,
        "temperature": temperature
    }

    should_close_session = False
    if session is None:
        session = aiohttp.ClientSession()
        should_close_session = True

    try:
        for attempt in range(max_retries):
            try:
                async with session.post(API_URL, json=data, headers=headers) as response:
                    if response.status == 200:
                        resp_json = await response.json()
                        if 'choices' in resp_json:
                            logger.info(
                                "Successfully received response from OpenRouter")
                            return resp_json['choices'][0]['message']['content']
                    elif response.status == 429:  # Rate limit exceeded
                        delay = base_delay * (2 ** attempt)
                        logger.warning(
                            f"Rate limit exceeded. Retrying in {delay} seconds...")
                        await asyncio.sleep(delay)
                        continue
                    else:
                        error_text = await response.text()
                        logger.error(
                            f"Error response (status {response.status}): {error_text}")

                    delay = base_delay * (2 ** attempt)
                    logger.warning(
                        f"Attempt {attempt+1} failed. Retrying in {delay} seconds...")
                    await asyncio.sleep(delay)

            except Exception as e:
                if attempt == max_retries - 1:
                    logger.error(
                        f"All retry attempts failed. Final error: {str(e)}")
                    raise
                delay = base_delay * (2 ** attempt)
                logger.warning(
                    f"Error occurred: {str(e)}. Retrying in {delay} seconds...")
                await asyncio.sleep(delay)
    finally:
        if should_close_session:
            await session.close()

    raise Exception("Max retries exceeded")


async def analyze_email_needs_response(email_data: Dict, session: aiohttp.ClientSession) -> bool:
    """Analyze a single email to determine if it needs a response."""
    logger.info(f"Analyzing email: {email_data['subject']}")

    # Log available data
    logger.info(f"Email data keys: {email_data.keys()}")
    logger.info(f"Thread content present: {'thread_content' in email_data}")

    email_content = f"""
From: {email_data['from']}
Subject: {email_data['subject']}
Content: {email_data.get('thread_content', email_data.get('snippet', '[No content]'))}
    """

    try:
        response = await make_api_call(
            prompt=f"{email_response_analysis}\n\nEmail to analyze:\n{email_content}",
            session=session
        )

        logger.info(f"Got response for {email_data['subject']}")
        logger.debug(f"Response content: {response}")

        # Parse the response...
        start_tag = "<NEEDS_RESPONSE>"
        end_tag = "</NEEDS_RESPONSE>"
        start_idx = response.find(start_tag) + len(start_tag)
        end_idx = response.find(end_tag)

        if start_idx != -1 and end_idx != -1:
            needs_response = response[start_idx:end_idx].strip(
            ).lower() == "true"
            logger.info(
                f"Analysis result for {email_data['subject']}: {needs_response}")
            return needs_response
        return False

    except Exception as e:
        logger.error(f"Error analyzing email: {e}")
        return False


async def get_emails_needing_response(service) -> List[Dict]:
    """Retrieve new emails and analyze which ones need responses."""
    try:
        # Initialize indexer and get new emails
        indexer = GmailIndexer(service)
        new_emails = indexer.get_new_emails()

        if not new_emails:
            logger.info("No new emails found to analyze")
            return []

        # Get full message content for each email
        emails_with_content = []
        for email in new_emails:
            # Get the full thread content
            thread = service.users().threads().get(
                userId='me',
                id=email['threadId'],
                format='full'
            ).execute()

            # Extract thread content
            thread_content = []
            for msg in thread['messages']:
                headers = msg['payload']['headers']
                from_header = next(
                    (h['value'] for h in headers if h['name'].lower() == 'from'), '')
                date = next((h['value']
                            for h in headers if h['name'].lower() == 'date'), '')

                # Extract message body
                body = ""
                if 'parts' in msg['payload']:
                    for part in msg['payload']['parts']:
                        mime_type = part.get('mimeType')
                        if mime_type in ('text/plain', 'text/html'):
                            part_body = part.get('body', {}).get('data', '')
                            if part_body:
                                decoded_body = base64.urlsafe_b64decode(
                                    part_body).decode('utf-8', errors='replace')
                                if decoded_body.strip():
                                    body = decoded_body.strip()
                                    break
                elif 'body' in msg['payload'] and 'data' in msg['payload']['body']:
                    body = base64.urlsafe_b64decode(
                        msg['payload']['body']['data']
                    ).decode('utf-8', errors='replace')

                if body.strip():
                    thread_content.extend([
                        f"Date: {date}",
                        f"From: {from_header}",
                        "",
                        body.strip(),
                        ""
                    ])

            # Add thread content to email data
            email['thread_content'] = "\n".join(thread_content)
            logger.info(f"Processed email: {email['subject']}")
            logger.info(
                f"Thread content length: {len(email['thread_content'])}")
            emails_with_content.append(email)

        emails_needing_response = []
        async with aiohttp.ClientSession() as session:
            tasks = [
                analyze_email_needs_response(email, session)
                for email in emails_with_content
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)

            for email, needs_response in zip(emails_with_content, results):
                if isinstance(needs_response, Exception):
                    logger.error(f"Error analyzing email: {needs_response}")
                    continue

                if needs_response:
                    emails_needing_response.append(email)

        logger.info(
            f"Found {len(emails_needing_response)} emails needing response out of {len(new_emails)} new emails")
        return emails_needing_response

    except Exception as e:
        logger.error(f"Error in get_emails_needing_response: {e}")
        logger.exception(e)  # This will print the full stack trace
        return []

# Example usage function


async def analyze_new_emails():
    """Helper function to run the email analysis"""
    from ..GCP.gmail_auth import get_gmail_service

    try:
        service = get_gmail_service()
        emails_needing_response = await get_emails_needing_response(service)

        if emails_needing_response:
            logger.info("\nEmails needing response:")
            for email in emails_needing_response:
                logger.info(f"\nFrom: {email['from']}")
                logger.info(f"Subject: {email['subject']}")
                logger.info(f"Date: {email['date']}")
        else:
            logger.info("No emails requiring response found")

    except Exception as e:
        logger.error(f"Error in analyze_new_emails: {e}")

if __name__ == "__main__":
    asyncio.run(analyze_new_emails())


async def modify_draft_prompt(current_prompt: str, new_requirements: str, session: aiohttp.ClientSession = None) -> str:
    """Modify the existing draft prompt with new requirements."""
    try:
        prompt = draft_prompt_modification.format(
            current_prompt=current_prompt,
            new_requirements=new_requirements
        )

        response = await make_api_call(
            prompt=prompt,
            model="openai/o3-mini",
            temperature=0.3,
            session=session
        )

        # Extract the modified prompt from the response
        start_tag = "<MODIFIED_PROMPT>"
        end_tag = "</MODIFIED_PROMPT>"
        start_idx = response.find(start_tag) + len(start_tag)
        end_idx = response.find(end_tag)

        if start_idx != -1 and end_idx != -1:
            modified_prompt = response[start_idx:end_idx].strip()
            return modified_prompt

        logger.error("Failed to parse modified prompt from response")
        return current_prompt

    except Exception as e:
        logger.error(f"Error modifying draft prompt: {e}")
        return current_prompt


async def analyze_response_requirements(email_data: Dict, session: aiohttp.ClientSession) -> Dict:
    """Analyze if an email can be auto-responded to or needs user input."""
    # Get full email thread content
    thread_content = email_data.get('thread_content', '')

    # Ensure we have the full message content
    if not thread_content:
        logger.warning("No thread content found in email_data")
        return {
            "can_auto_respond": False,
            "required_info": [],
            "search_terms": [],
            "reasoning": "Error: No email content found to analyze"
        }

    email_content = f"""
From: {email_data['from']}
Subject: {email_data['subject']}
Thread Content:
{thread_content}
    """

    try:
        response = await make_api_call(
            prompt=f"{email_response_requirements_analysis}\n\nEmail to analyze:\n{email_content}",
            model="openai/o3-mini",
            temperature=0.3,
            session=session
        )

        # Parse response
        response_type_start = response.find(
            "<RESPONSE_TYPE>") + len("<RESPONSE_TYPE>")
        response_type_end = response.find("</RESPONSE_TYPE>")
        response_type = response[response_type_start:response_type_end].strip()

        required_info_start = response.find(
            "<REQUIRED_INFO>") + len("<REQUIRED_INFO>")
        required_info_end = response.find("</REQUIRED_INFO>")
        required_info = response[required_info_start:required_info_end].strip()

        # Parse search terms as a list
        search_terms_start = response.find(
            "<SEARCH_TERMS>") + len("<SEARCH_TERMS>")
        search_terms_end = response.find("</SEARCH_TERMS>")
        search_terms_str = response[search_terms_start:search_terms_end].strip(
        )

        # Extract terms from array format
        try:
            search_terms = ast.literal_eval(
                search_terms_str) if search_terms_str else []
        except:
            search_terms = []
            logger.error("Failed to parse search terms array")

        reasoning_start = response.find("Reasoning: ") + len("Reasoning: ")
        reasoning_end = response.find("\n\nDecision:")
        reasoning = response[reasoning_start:reasoning_end].strip()

        return {
            "can_auto_respond": response_type == "AUTO_RESPOND",
            "required_info": [info.strip() for info in required_info.split('\n') if info.strip()],
            "search_terms": search_terms,
            "reasoning": reasoning
        }

    except Exception as e:
        logger.error(f"Error in response requirements analysis: {e}")
        return {
            "can_auto_respond": False,
            "required_info": [],
            "search_terms": [],
            "reasoning": f"Error in analysis: {str(e)}"
        }


async def gather_context_and_create_draft(email_data: Dict, search_terms: List[str], service, session: aiohttp.ClientSession) -> bool:
    try:
        # Get relevant context from previous emails
        all_context = []
        for term in search_terms:
            search_result = search_emails(
                service=service,
                keywords=term,
                newer_than_days=30
            )
            if search_result.strip():  # Check if we got any content
                all_context.append(search_result)

        combined_thread = f"""
Original Thread:
{email_data.get('thread_content', '')}

Relevant Context from Previous Emails:
{'\n\n'.join(all_context)}
"""

        with open("src/gmail_mcp/utils/constants/draft_prompt.txt", "r") as f:
            draft_prompt = f.read()

        response = await make_api_call(
            prompt=draft_response_generation.format(
                draft_prompt=draft_prompt,
                email_thread=combined_thread
            ),
            model="anthropic/claude-3-sonnet",
            temperature=0.7,
            session=session
        )

        start_tag = "<DRAFT_RESPONSE>"
        end_tag = "</DRAFT_RESPONSE>"
        start_idx = response.find(start_tag) + len(start_tag)
        end_idx = response.find(end_tag)
        draft_content = response[start_idx:end_idx].strip()

        # Create draft in Gmail
        draft = {
            'message': {
                'raw': base64.urlsafe_b64encode(draft_content.encode('utf-8')).decode('utf-8'),
                'threadId': email_data['threadId']
            }
        }
        service.users().drafts().create(userId='me', body=draft).execute()

        logger.info(f"Created draft response for: {email_data['subject']}")
        return True

    except Exception as e:
        logger.error(f"Error creating draft response: {e}")
        return False
