#!/usr/bin/env python3
"""
Email Responder Bot - Automatyczne odpowiedzi z LLM (refactored slim version)
Deleguje logikę do modułów w llmass/*
"""
import os
import sys
import argparse
import email
from datetime import datetime, timedelta
from typing import Dict, List
from dotenv import load_dotenv

# Import modułów llmass
from llmass.imap import connect_imap
from llmass.organizer.parser import get_email_content, detect_imap_server
from llmass.responder.drafts import detect_smtp_server, resolve_drafts_folder_name, save_draft, is_auto_reply
from llmass.responder.conversation import fetch_conversation_history, detect_language, get_greeting

# Conditional imports dla LLM
try:
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    import torch
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    AutoTokenizer = None
    AutoModelForCausalLM = None
    BitsAndBytesConfig = None
    torch = None
    TRANSFORMERS_AVAILABLE = False


class EmailResponder:
    """Email Responder Bot - generuje odpowiedzi na emaile używając LLM."""
    
    def __init__(self, email_address: str, password: str, 
                 model_name: str = "Qwen/Qwen2.5-7B-Instruct",
                 imap_server: str = None, smtp_server: str = None):
        """Inicjalizacja bota."""
        self.email_address = email_address
        self.password = password
        self.imap_server = imap_server or detect_imap_server(email_address)
        self.smtp_server = smtp_server or detect_smtp_server(email_address)
        self.model_name = model_name
        self.imap = None
        self.client = None
        self.model = None
        self.tokenizer = None
        
        # Generation params
        self.generation_params = {
            'max_new_tokens': int(os.getenv('MAX_TOKENS', '500')),
            'temperature': float(os.getenv('TEMPERATURE', '0.7')),
            'top_p': 0.9,
            'do_sample': True,
            'pad_token_id': None,  # Will be set after tokenizer load
        }
        
        # Signature (bez pozdrowienia - będzie dodane dynamicznie)
        self.sender_name = os.getenv('SENDER_NAME', 'Bot')
        self.sender_title = os.getenv('SENDER_TITLE', '')
        self.sender_company = os.getenv('SENDER_COMPANY', '')
        signature_template = os.getenv('SIGNATURE', '{SENDER_NAME}\n{SENDER_TITLE}\n{SENDER_COMPANY}')
        # Replace literal \n with actual newlines
        signature_template = signature_template.replace('\\n', '\n')
        self.signature = signature_template.format(
            SENDER_NAME=self.sender_name,
            SENDER_TITLE=self.sender_title,
            SENDER_COMPANY=self.sender_company
        )
        
        # Prompt template
        self.system_prompt = f"""Jesteś pomocnym asystentem emailowym. Twoje zadanie to pisanie profesjonalnych, zwięzłych i pomocnych odpowiedzi na emaile.

INSTRUKCJE:
1. Odpowiadaj w języku emaila wejściowego
2. Zachowaj profesjonalny i uprzejmy ton
3. Bądź konkretny i pomocny
4. Jeśli email zawiera pytanie, odpowiedz bezpośrednio
5. Jeśli to newsletter/powiadomienie, podziękuj i potwierdź otrzymanie

PODPIS:
{self.signature}
"""
        
        self.prompt_template = """EMAIL ORYGINALNY:
Od: {from}
Temat: {subject}

{body}

NAPISZ ODPOWIEDŹ:"""
    
    def load_model(self):
        """Ładuje model LLM."""
        if not TRANSFORMERS_AVAILABLE:
            print("⚠️  Transformers nie jest dostępne - używam trybu mock")
            return False
        
        try:
            print(f"🤖 Ładowanie modelu {self.model_name}...")
            
            device = os.getenv('DEVICE', 'cpu').lower()
            if device == 'gpu' and torch and torch.cuda.is_available():
                print(f"   Używam urządzenia: cuda")
                
                # 4-bit quantization config
                quant_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_compute_dtype=torch.float16,
                    bnb_4bit_use_double_quant=True,
                    bnb_4bit_quant_type="nf4"
                )
                
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name,
                    quantization_config=quant_config,
                    device_map="auto",
                    trust_remote_code=True
                )
            else:
                print(f"   Używam urządzenia: cpu")
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name,
                    torch_dtype=torch.float32,
                    device_map="cpu",
                    trust_remote_code=True
                )
            
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
            
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            
            self.generation_params['pad_token_id'] = self.tokenizer.pad_token_id
            
            print("✅ Model załadowany pomyślnie!")
            return True
            
        except Exception as e:
            print(f"❌ Błąd ładowania modelu: {e}")
            self.model = None
            return False
    
    def connect(self):
        """Połączenie z serwerem IMAP."""
        try:
            imap_retries = int(os.getenv('IMAP_RETRIES', '2'))
            imap_backoff = float(os.getenv('IMAP_BACKOFF', '0.5'))
            self.client, self.imap = connect_imap(self.imap_server, self.email_address, self.password,
                                                   ssl=True, retries=imap_retries, backoff=imap_backoff, verbose=False)
            print(f"✅ Połączono z {self.imap_server}")
            return True
        except Exception as e:
            print(f"❌ Błąd połączenia: {e}")
            return False
    
    def disconnect(self):
        """Rozłączenie z serwerem."""
        if self.imap:
            self.imap.close()
            self.imap.logout()
            print("👋 Rozłączono z serwerem")
    
    def generate_response_with_llm(self, email_content: Dict) -> str:
        """Generuje odpowiedź używając modelu LLM."""
        if not TRANSFORMERS_AVAILABLE or not self.model:
            return self._generate_mock_response(email_content)
        
        try:
            # Fetch full conversation history
            sender = email_content.get('from', '')
            sent_history, received_history = fetch_conversation_history(self, sender, limit=3)
            
            # Build context with full conversation
            context = ""
            if sent_history or received_history:
                context = "\n\nHISTORIA KONWERSACJI:\n"
                
                # Combine and sort by date (interleave sent/received)
                all_history = []
                for h in sent_history:
                    all_history.append(('WYSŁANE', h))
                for h in received_history:
                    all_history.append(('OTRZYMANE', h))
                
                # Simple ordering by position (already sorted newest first in each list)
                for typ, h in all_history[:5]:  # Limit total to 5
                    context += f"\n[{typ}] {h.get('subject', 'Brak tematu')}\n{h.get('body', '')[:200]}...\n"
            
            # Build prompt
            prompt_content = self.prompt_template.format(
                **{
                    'from': email_content.get('from', 'Nieznany'),
                    'subject': email_content.get('subject', 'Brak tematu'),
                    'body': email_content.get('body', '')[:1000]  # Limit
                }
            )
            
            messages = [
                {"role": "system", "content": self.system_prompt + context},
                {"role": "user", "content": prompt_content}
            ]
            
            text = self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
            
            model_inputs = self.tokenizer([text], return_tensors="pt")
            
            # Move to correct device
            if torch and torch.cuda.is_available() and next(self.model.parameters()).is_cuda:
                model_inputs = model_inputs.to('cuda')
            
            # Generate
            try:
                generated_ids = self.model.generate(
                    model_inputs.input_ids,
                    **self.generation_params
                )
            except RuntimeError as e:
                if "out of memory" in str(e).lower():
                    print("❗ CUDA OOM podczas generowania. Używam mock response...")
                    if torch and torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    return self._generate_mock_response(email_content)
                raise
            
            generated_ids = [
                output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
            ]
            
            response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
            return response.strip()
            
        except Exception as e:
            print(f"❌ Błąd generowania: {e}")
            return self._generate_mock_response(email_content)
    
    def _generate_mock_response(self, email_content: Dict) -> str:
        """Generuje przykładową odpowiedź gdy model nie jest dostępny."""
        subject = email_content.get('subject', 'Brak tematu')
        body = email_content.get('body', '')
        
        # Detect language
        text = f"{subject} {body}"
        lang = detect_language(text)
        
        # Fetch conversation history
        sender = email_content.get('from', '')
        sent_history, received_history = fetch_conversation_history(self, sender, limit=2)
        
        # Build history context
        history_text = ""
        if sent_history or received_history:
            history_count = len(sent_history) + len(received_history)
            if lang == 'pl':
                history_text = f"\n\nNawiązuję do naszej wcześniejszej korespondencji ({history_count} wiadomości).\n"
            elif lang == 'de':
                history_text = f"\n\nIch beziehe mich auf unsere frühere Korrespondenz ({history_count} Nachrichten).\n"
            elif lang == 'fr':
                history_text = f"\n\nJe fais référence à notre correspondance précédente ({history_count} messages).\n"
            else:  # en and others
                history_text = f"\n\nReferring to our previous correspondence ({history_count} messages).\n"
        
        # Language-specific templates
        templates = {
            'pl': f"""Dziękuję za Twoją wiadomość dotyczącą "{subject}".{history_text}
Przeanalizowałem treść Twojego emaila i chętnie pomogę w tej sprawie.
Twoje zapytanie jest dla mnie ważne i postaram się odpowiedzieć jak najszybciej.

Jeśli potrzebujesz dodatkowych informacji, proszę daj mi znać.""",
            
            'de': f"""Vielen Dank für Ihre Nachricht bezüglich "{subject}".{history_text}
Ich habe den Inhalt Ihrer E-Mail analysiert und helfe Ihnen gerne in dieser Angelegenheit.
Ihre Anfrage ist mir wichtig und ich werde mich bemühen, so schnell wie möglich zu antworten.

Wenn Sie zusätzliche Informationen benötigen, lassen Sie es mich bitte wissen.""",
            
            'fr': f"""Merci pour votre message concernant "{subject}".{history_text}
J'ai analysé le contenu de votre email et je serais heureux de vous aider dans cette affaire.
Votre demande est importante pour moi et je m'efforcerai de répondre le plus rapidement possible.

Si vous avez besoin d'informations supplémentaires, n'hésitez pas à me le faire savoir.""",
            
            'en': f"""Thank you for your message regarding "{subject}".{history_text}
I have analyzed the content of your email and I would be happy to help you with this matter.
Your inquiry is important to me and I will try to respond as quickly as possible.

If you need any additional information, please let me know.""",
        }
        
        response_body = templates.get(lang, templates['en'])
        greeting = get_greeting(lang)
        
        response = f"""{response_body}

{greeting}
{self.signature}

---
[Ta odpowiedź została wygenerowana automatycznie przez Email Responder Bot]"""
        
        return response
    
    def set_generation_params(self, **kwargs):
        """Ustawia parametry generowania odpowiedzi."""
        self.generation_params.update(kwargs)
        print(f"📝 Zaktualizowano parametry generowania: {kwargs}")
    
    def process_emails(self, folder: str = "INBOX", limit: int = 100,
                      filter_unread: bool = True, dry_run: bool = False,
                      since_days: int = 7, since_date: str = None):
        """Przetwarza emaile i generuje odpowiedzi."""
        print(f"\n📧 Przetwarzam emaile z folderu: {folder}\n")
        
        # Select folder
        result, data = self.client.safe_select(folder, readonly=False)
        if result != 'OK':
            print(f"❌ Nie można otworzyć folderu {folder}")
            return
        
        # Search
        search_criteria = 'UNSEEN' if filter_unread else 'ALL'
        
        # Date filter
        if since_date:
            cutoff = datetime.strptime(since_date, '%Y-%m-%d')
            search_str = cutoff.strftime('%d-%b-%Y')
            result, data = self.client.safe_uid('SEARCH', None, f'{search_criteria} SINCE {search_str}')
        elif since_days:
            cutoff = datetime.now() - timedelta(days=since_days)
            search_str = cutoff.strftime('%d-%b-%Y')
            result, data = self.client.safe_uid('SEARCH', None, f'{search_criteria} SINCE {search_str}')
        else:
            result, data = self.client.safe_uid('SEARCH', None, search_criteria)
        
        if result != 'OK' or not data or not data[0]:
            print("📭 Brak emaili do przetworzenia")
            return
        
        email_ids = data[0].split()
        if not email_ids:
            print("📭 Brak emaili do przetworzenia")
            return
        
        # Limit
        email_ids = email_ids[-limit:] if len(email_ids) > limit else email_ids
        email_ids.reverse()  # Newest first
        
        print(f"⏱️  Filtr czasu: od {cutoff.strftime('%d-%b-%Y') if 'cutoff' in locals() else 'początku'}, limit: {limit}")
        print(f"📊 Znaleziono {len(email_ids)} emaili do przetworzenia\n")
        
        processed = 0
        drafts_created = 0
        
        for idx, email_id in enumerate(email_ids, 1):
            try:
                print(f"--- Email {idx}/{len(email_ids)} ---")
                
                # Fetch
                result, data = self.client.safe_uid('FETCH', email_id, '(RFC822)')
                if result != 'OK' or not data or not data[0]:
                    print("❌ Błąd pobierania emaila\n")
                    continue
                
                raw_email = data[0][1]
                msg = email.message_from_bytes(raw_email)
                email_content = get_email_content(msg)
                
                # Skip auto-replies
                if is_auto_reply(email_content):
                    print("⏭️  Pomijam automatyczną odpowiedź\n")
                    continue
                
                print(f"📩 Od: {email_content.get('from', 'Nieznany')[:50]}")
                print(f"📋 Temat: {email_content.get('subject', 'Brak tematu')[:50]}")
                
                # Generate response
                print("🤖 Generuję odpowiedź...")
                response = self.generate_response_with_llm(email_content)
                
                print(f"\n📝 Wygenerowana odpowiedź:")
                print("-" * 50)
                print(response[:500] + ("..." if len(response) > 500 else ""))
                print("-" * 50)
                
                # Save draft
                if not dry_run:
                    if save_draft(self, email_content, response):
                        drafts_created += 1
                else:
                    print("🧪 [DRY-RUN] Draft nie został zapisany")
                
                processed += 1
                print()
                
            except KeyboardInterrupt:
                print("\n⚠️  Przerwano przez użytkownika")
                break
            except Exception as e:
                print(f"❌ Błąd przetwarzania emaila: {e}\n")
                continue
        
        print(f"📊 Podsumowanie:")
        print(f"   - Przetworzono: {processed} emaili")
        print(f"   - Utworzono draftów: {drafts_created}")


def main():
    load_dotenv()
    parser = argparse.ArgumentParser(description='Email Responder Bot z LLM')
    parser.add_argument('--email', help='Adres email', default=os.getenv('EMAIL_ADDRESS'))
    parser.add_argument('--password', help='Hasło', default=os.getenv('EMAIL_PASSWORD'))
    parser.add_argument('--imap-server', help='Serwer IMAP', default=os.getenv('IMAP_SERVER'))
    parser.add_argument('--smtp-server', help='Serwer SMTP', default=os.getenv('SMTP_SERVER'))
    parser.add_argument('--model', help='Model LLM', default=os.getenv('MODEL_NAME', 'Qwen/Qwen2.5-7B-Instruct'))
    parser.add_argument('--folder', default='INBOX', help='Folder emaili')
    parser.add_argument('--limit', type=int, default=int(os.getenv('LIMIT', '10')), help='Limit emaili')
    parser.add_argument('--all', action='store_true', help='Przetwarzaj wszystkie (nie tylko nieprzeczytane)')
    parser.add_argument('--dry-run', action='store_true', help='Tylko generuj, nie zapisuj draftów')
    parser.add_argument('--since-days', type=int, default=int(os.getenv('SINCE_DAYS', '7')), help='Ile dni wstecz')
    parser.add_argument('--since-date', help='Data (YYYY-MM-DD)')
    parser.add_argument('--offline', action='store_true', help='Tryb offline (mock responses)')
    parser.add_argument('--temperature', type=float, help='Temperature dla LLM')
    parser.add_argument('--max-tokens', type=int, help='Max tokens dla LLM')
    
    args = parser.parse_args()
    
    if not args.email or not args.password:
        print("❌ Brak danych logowania. Podaj --email/--password lub .env")
        sys.exit(1)
    
    bot = EmailResponder(args.email, args.password, args.model, args.imap_server, args.smtp_server)
    
    if args.temperature:
        bot.set_generation_params(temperature=args.temperature)
    if args.max_tokens:
        bot.set_generation_params(max_new_tokens=args.max_tokens)
    
    if not args.offline:
        if not bot.load_model():
            print("⚠️  Kontynuuję w trybie mock (bez modelu)")
    else:
        print("📴 Tryb offline - używam mock responses")
    
    if bot.connect():
        try:
            bot.process_emails(
                folder=args.folder,
                limit=args.limit,
                filter_unread=not args.all,
                dry_run=args.dry_run,
                since_days=args.since_days,
                since_date=args.since_date
            )
        finally:
            bot.disconnect()


if __name__ == "__main__":
    main()
