import imaplib
import email
from email.header import decode_header
from datetime import datetime
import uuid
from app import mysql
from app.utils.company_settings import get_company_smtp_settings
from app.utils.celery_app import celery
import logging
from app.utils.llm_reply import generate_llm_reply
from app.utils.email_sender import send_simple_email
from app.routes.api import send_email, schedule_meeting
import time
import socket
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

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

def connect_to_imap(imap_host, imap_port, email_username, email_password, max_retries=3):
    """Helper function to establish IMAP connection with retries"""
    for attempt in range(max_retries):
        try:
            mail = imaplib.IMAP4_SSL(imap_host, imap_port, timeout=30)
            mail.login(email_username, email_password)
            return mail
        except (socket.error, imaplib.IMAP4.error) as e:
            if attempt < max_retries - 1:
                logger.warning(f"IMAP connection attempt {attempt + 1} failed: {str(e)}. Retrying...")
                time.sleep(5)  # Wait 5 seconds before retrying
            else:
                raise
    return None

@celery.task
def scan_email_replies():
    """
    Scan company email inboxes for replies to campaign emails
    """
    logger.info("=== Starting email reply scan ===")
    logger.info(f"Scan started at: {datetime.now()}")
    
    cursor = mysql.connection.cursor()
    
    try:
        # Get all companies with SMTP settings
        cursor.execute("""
            SELECT id, mail_username, mail_password, mail_host, imap_host, imap_port, company_id, mail_port,mail_from_name, mail_from_address
            FROM company_settings
            WHERE mail_username IS NOT NULL
        """)
        companies = cursor.fetchall()
        
        logger.info(f"Found {len(companies)} companies to scan")
        
        for company in companies:
            company_id = company[6]
            email_username = company[1]
            email_password = company[2]
            smtp_host = company[3]
            smtp_port = company[7]
            imap_host = company[4]
            imap_port = company[5]
            mail_from_name = company[8]
            mail_from_address = company[9]
            
            logger.info(f"Processing company {company_id}")
            logger.info(f"Connecting to IMAP server: {imap_host}:{imap_port}")
            
            ai_responses_to_send = []
            mail = None
            try:
                mail = connect_to_imap(imap_host, imap_port, email_username, email_password)
                if not mail:
                    logger.error(f"Failed to connect to IMAP server for company {company_id}")
                    continue
                
                logger.info("IMAP SSL connection established")
                mail.select('INBOX')
                
                # Search for unread emails
                _, messages = mail.search(None, 'UNSEEN')
                message_count = len(messages[0].split())
                logger.info(f"Found {message_count} unread messages")
                
                for msg_num in messages[0].split():
                    logger.info(f"Processing message {msg_num}")
                    _, msg_data = mail.fetch(msg_num, '(RFC822)')
                    email_body = msg_data[0][1]
                    email_message = email.message_from_bytes(email_body)
                    
                    # Debug email headers
                    # logger.info("Email Headers:")
                    logger.info(f"From: {email_message.get('From')}")
                    logger.info(f"Subject: {email_message.get('Subject')}")
                    # logger.info(f"Message-ID: {email_message.get('Message-ID')}")
                    # logger.info(f"References: {email_message.get('References')}")
                    # logger.info(f"In-Reply-To: {email_message.get('In-Reply-To')}")
                    
                    # Get References and In-Reply-To headers
                    references = email_message.get('References')
                    in_reply_to = email_message.get('In-Reply-To')

                    # Collect all possible referenced message IDs
                    referenced_ids = []
                    if references:
                        # References can contain multiple message IDs separated by spaces or newlines
                        # referenced_ids += [msgid.strip(' \n\r<>') for msgid in references.replace('\n', ' ').split()]
                        referenced_ids += [msgid.strip(' \n\r') for msgid in references.replace('\n', ' ').split()]
                    if in_reply_to:
                        referenced_ids.append(in_reply_to.strip(' \n\r<>'))

                    logger.info(f"Looking for campaign with referenced message IDs: {referenced_ids}")

                    found = False
                    for ref_id in referenced_ids:
                        cursor.execute("""
                            SELECT cr.campaign_id, cr.contact_id, c.email
                            FROM campaign_recipients cr
                            JOIN contacts c ON cr.contact_id = c.id
                            WHERE cr.message_id = %s
                        """, (ref_id,))
                        result = cursor.fetchone()
                        if result:
                            campaign_id, contact_id, contact_email = result
                            logger.info(f"Found matching campaign: {campaign_id} for contact: {contact_email} (ref_id: {ref_id})")
                            found = True
                            break

                    if not found:
                        logger.info(f"No campaign found for any referenced message ID: {referenced_ids}")
                        continue
                    
                    # Extract reply content
                    reply_content = ""
                    if email_message.is_multipart():
                        for part in email_message.walk():
                            if part.get_content_type() == "text/plain":
                                try:
                                    reply_content = part.get_payload(decode=True).decode()
                                except:
                                    reply_content = part.get_payload()
                                break
                    else:
                        try:
                            reply_content = email_message.get_payload(decode=True).decode()
                        except:
                            reply_content = email_message.get_payload()
                    
                    logger.info(f"Extracted reply content: {reply_content[:100]}...")  # Log first 100 chars
                    
                    # Clean reply content
                    clean_reply = extract_clean_reply(reply_content)
                    logger.info(f"Cleaned reply content: {clean_reply[:100]}...")  # Log first 100 chars
                    
                    # Store the clean reply
                    reply_id = str(uuid.uuid4())
                    try:
                        cursor.execute("""
                            INSERT INTO email_replies (
                                id, campaign_id, contact_id, message_id,
                                reply_content, reply_date, created_at
                            ) VALUES (%s, %s, %s, %s, %s, %s, NOW())
                        """, (
                            reply_id,
                            campaign_id,
                            contact_id,
                            email_message.get('Message-ID', '').strip('<>'),
                            clean_reply,
                            datetime.strptime(email_message.get('Date'), '%a, %d %b %Y %H:%M:%S %z')
                        ))
                        
                        mysql.connection.commit()
                        message_id = email_message.get('Message-ID', '').strip('<>')
                        logger.info(f"Successfully stored reply for campaign {campaign_id}")
                    except Exception as db_error:
                        logger.error(f"Error storing reply: {str(db_error)}")
                        logger.error(f"SQL Parameters: campaign_id={campaign_id}, contact_id={contact_id}")
                        continue
                    
                    # Mark email as read
                    mail.store(msg_num, '+FLAGS', '\\Seen')
                    logger.info(f"Marked message {msg_num} as read")
                    
                    ai_responses_to_send.append({
                        "campaign_id": campaign_id,
                        "contact_id": contact_id,
                        "contact_email": contact_email,
                        "clean_reply": clean_reply,
                        "email_message": email_message,
                        "company_id": company_id,
                        "smtp_host": smtp_host,
                        "smtp_port": smtp_port,
                        "email_username": email_username,
                        "email_password": email_password,
                        "imap_port": imap_port,
                        "mail_from_name": mail_from_name,
                        "mail_from_address": mail_from_address,
                    })
                
                # Close IMAP connection as soon as possible
                mail.close()
                mail.logout()
            except (socket.error, imaplib.IMAP4.error) as e:
                logger.error(f"IMAP connection error for company {company_id}: {str(e)}")
                time.sleep(30)
                continue
            except Exception as e:
                logger.error(f"Error processing company {company_id}: {str(e)}")
                continue
            finally:
                try:
                    mail.close()
                    mail.logout()
                except:
                    pass
            
            # === OUTSIDE IMAP: Now send AI responses ===
            for item in ai_responses_to_send:
                try:
                    llm_result = generate_llm_reply(
                        item["company_id"],
                        item["contact_email"],
                        item["contact_email"],
                        item["clean_reply"],
                        campaign_id=item["campaign_id"],
                        contact_id=item["contact_id"]
                    )
                    llm_reply = llm_result["response"]
                    meeting_action = llm_result["action"]

                    cursor.execute("SELECT subject FROM email_campaigns WHERE id = %s", (item["campaign_id"],))
                    campaign_subject = cursor.fetchone()[0]
                    reply_subject = f"Re: {campaign_subject}"

                    cursor.execute("""
                        SELECT id FROM contacts 
                        WHERE email = %s AND company_id = %s
                    """, (item["contact_email"], item["company_id"]))
                    contact_result = cursor.fetchone()
                    CONTACT_ID = contact_result[0] if contact_result else None

                    cursor.execute("""
                    SELECT name,email FROM companies 
                    WHERE id = %s
                    """, (item["company_id"],))
                    comapnyResult = cursor.fetchone()
                    companyName = comapnyResult[0] if comapnyResult else None
                    companyEmail = comapnyResult[1] if comapnyResult else None

                    if meeting_action and meeting_action.get("action") == "schedule_meeting":
                        # Get the original message_id from the campaign
                        cursor.execute("""
                            SELECT message_id 
                            FROM campaign_recipients 
                            WHERE campaign_id = %s AND contact_id = %s
                        """, (item["campaign_id"], item["contact_id"]))
                        result = cursor.fetchone()
                        original_message_id = result[0] if result else None
                        
                        if original_message_id:
                            # Ensure message_id has <> signs
                            if not original_message_id.startswith('<'):
                                original_message_id = f"<{original_message_id}"
                            if not original_message_id.endswith('>'):
                                original_message_id = f"{original_message_id}>"
                        
                        # Call your schedule_meeting logic (reuse your API logic)
                        meeting_response = schedule_meeting(item["company_id"], meeting_action)
                        logger.info(f"inside meeting if")
                        
                        # Send the meeting email with Google Calendar button
                        send_email(
                            to_email=item["contact_email"],
                            subject="Meeting Scheduled Confirmation",
                            date=meeting_action["date"],
                            time=meeting_action["time"],
                            company_name=companyName,  # Or fetch from DB
                            company_email=companyEmail,      # Or fetch from DB
                            customer_email=item["contact_email"],
                            company_id=item["company_id"],
                            message_id=original_message_id,    # Pass the original message_id
                            references=original_message_id,    # Pass the original message_id as references
                            in_reply_to=original_message_id    # Pass the original message_id as in-reply-to
                        )
                        
                        # Store the AI's meeting confirmation response
                        reply_id = str(uuid.uuid4())
                        cursor.execute("""
                            INSERT INTO email_replies (
                                id, campaign_id, contact_id, message_id,
                                reply_content, reply_date, created_at,
                                is_ai_response
                            ) VALUES (%s, %s, %s, %s, %s, %s, NOW(), %s)
                        """, (
                            reply_id,
                            item["campaign_id"],
                            CONTACT_ID,
                            original_message_id,  # Use the original message_id instead of None
                            "Your meeting has been scheduled. Please check your email for details.",
                            datetime.now(),
                            True  # Mark as AI response
                        ))
                        mysql.connection.commit()
                    else:
                        # Normal LLM reply: send a simple email
                        contact_message_id = item["email_message"].get('Message-ID')
                        references_header = item["email_message"].get('References', '')
                        if references_header:
                            # Append the current message id to the references chain
                            references = f"{references_header} {contact_message_id}".strip()
                        else:
                            references = contact_message_id

                        # Generate a new unique Message-ID for the AI response
                        ai_message_id = f"<{str(uuid.uuid4())}@{item['smtp_host']}>"
                       
                        msg = MIMEMultipart()
                        msg['From'] = f"{mail_from_name} <{mail_from_address}>"
                        msg['To'] = item["contact_email"]
                        msg['Subject'] = reply_subject
                        msg['Message-ID'] = ai_message_id
                        msg['In-Reply-To'] = contact_message_id
                        msg['References'] = references
                        msg.attach(MIMEText(llm_reply, 'plain'))

                        # Send the email using SMTP
                        with smtplib.SMTP_SSL(item['smtp_host'], int(item['smtp_port']), timeout=30) as smtp:
                            smtp.login(item['email_username'], item['email_password'])
                            smtp.send_message(msg)
                        logger.info(f"Email reply from Ai")
                        
                        # Store the AI's response in email_replies with the new ai_message_id
                        reply_id = str(uuid.uuid4())
                        try:
                            cursor.execute("""
                                INSERT INTO email_replies (
                                    id, campaign_id, contact_id, message_id,
                                    reply_content, reply_date, created_at,
                                    is_ai_response
                                ) VALUES (%s, %s, %s, %s, %s, %s, NOW(), %s)
                            """, (
                                reply_id,
                                item["campaign_id"],
                                CONTACT_ID,
                                ai_message_id,  # Store the new AI message_id
                                llm_reply,
                                datetime.now(),
                                True
                            ))
                            mysql.connection.commit()
                            logger.info(f"Successfully stored AI response for campaign {item['campaign_id']}")
                        except Exception as db_error:
                            logger.error(f"Error storing AI response: {str(db_error)}")
                except Exception as e:
                    logger.error(f"Error sending AI response for campaign {item['campaign_id']}: {str(e)}")
            
            # Add a small delay between companies
            time.sleep(2)
        
    except Exception as e:
        logger.error(f"Error in scan_email_replies: {str(e)}")
    finally:
        cursor.close()
        logger.info("=== Finished email reply scan ===")

@celery.on_after_configure.connect
def setup_reply_scanner(sender, **kwargs):
    # Add new task to check for replies every 5 minutes
    sender.add_periodic_task(300.0, scan_email_replies.s(), name='Scan for email replies every 5 minutes')

def extract_clean_reply(body):
    """
    Extract only the new reply content from an email body.
    """
    # Common separators for quoted text in replies
    separators = [
        "\nOn ",  # Gmail, Outlook, etc.
        "\r\nOn ",
        "\nFrom:",
        "\r\nFrom:",
        "\n> ",  # Some clients use > for quoted lines
        "\r\n> ",
        "\n-----Original Message-----",
        "\r\n-----Original Message-----"
    ]
    for sep in separators:
        if sep in body:
            return body.split(sep)[0].strip()
    return body.strip()