"""
Ollama client for AI-powered article summarization
"""
import requests
import time
from datetime import datetime


class OllamaClient:
    """Client for communicating with Ollama server for text summarization"""
    
    def __init__(self, base_url, model, api_key=None, enabled=True, timeout=30):
        """
        Initialize Ollama client
        
        Args:
            base_url: Ollama server URL (e.g., http://localhost:11434)
            model: Model name to use (e.g., phi3:latest)
            api_key: Optional API key for authentication
            enabled: Whether Ollama is enabled
            timeout: Request timeout in seconds (default 30)
        """
        self.base_url = base_url.rstrip('/')
        self.model = model
        self.api_key = api_key
        self.enabled = enabled
        self.timeout = timeout

    def _chat_request(self, messages, options=None):
        """
        Helper to make chat requests to Ollama
        
        Args:
            messages: List of message dicts [{'role': 'user', 'content': '...'}]
            options: Optional dict of model parameters
            
        Returns:
            str: Generated text content
        """
        if options is None:
            options = {}
            
        url = f"{self.base_url}/api/chat"
        headers = {'Content-Type': 'application/json'}
        if self.api_key:
            headers['Authorization'] = f'Bearer {self.api_key}'
            
        payload = {
            'model': self.model,
            'messages': messages,
            'stream': False,
            'options': options
        }
        
        response = requests.post(
            url,
            json=payload,
            headers=headers,
            timeout=self.timeout
        )
        response.raise_for_status()
        
        result = response.json()
        return result.get('message', {}).get('content', '').strip()
        
    def summarize_article(self, content, max_words=150):
        """
        Summarize article content using Ollama
        
        Args:
            content: Full article text
            max_words: Maximum words in summary (default 150)
            
        Returns:
            {
                'summary': str,              # AI-generated summary
                'summary_word_count': int,   # Summary word count
                'original_word_count': int,  # Original article word count
                'success': bool,             # Whether summarization succeeded
                'error': str or None,        # Error message if failed
                'duration': float            # Time taken in seconds
            }
        """
        if not self.enabled:
            return {
                'summary': None,
                'summary_word_count': 0,
                'original_word_count': 0,
                'success': False,
                'error': 'Ollama is not enabled',
                'duration': 0
            }
        
        if not content or len(content.strip()) == 0:
            return {
                'summary': None,
                'summary_word_count': 0,
                'original_word_count': 0,
                'success': False,
                'error': 'Content is empty',
                'duration': 0
            }
        
        # Calculate original word count
        original_word_count = len(content.split())
        
        start_time = time.time()
        
        try:
            # Construct messages for chat API
            messages = [
                {
                    'role': 'system',
                    'content': f"You are a skilled journalist writing for The New York Times. Summarize the provided article in English in {max_words} words or less.\\n\\nWrite in the clear, engaging, and authoritative style of New York Times Magazine:\\n- Lead with the most newsworthy information\\n- Use active voice and vivid language\\n- Make it accessible and easy to read\\n- Focus on what matters to readers\\n- Even if the source is in German or another language, write your summary entirely in English\\n\\nIMPORTANT: Write in plain text only. Do NOT use markdown formatting (no ##, **, *, bullets, etc.). Just write natural prose."
                },
                {
                    'role': 'user',
                    'content': f"Summarize this article:\\n\\n{content}"
                }
            ]
            
            # Make request using chat endpoint
            summary = self._chat_request(
                messages, 
                options={
                    'temperature': 0.5,
                    'num_predict': 350
                }
            )
            
            if not summary:
                return {
                    'summary': None,
                    'summary_word_count': 0,
                    'original_word_count': original_word_count,
                    'success': False,
                    'error': 'Ollama returned empty summary',
                    'duration': time.time() - start_time
                }
            
            # Clean markdown formatting from summary
            summary = self._clean_markdown(summary)
            
            summary_word_count = len(summary.split())
            
            return {
                'summary': summary,
                'summary_word_count': summary_word_count,
                'original_word_count': original_word_count,
                'success': True,
                'error': None,
                'duration': time.time() - start_time
            }
            
        except requests.exceptions.Timeout:
            return {
                'summary': None,
                'summary_word_count': 0,
                'original_word_count': original_word_count,
                'success': False,
                'error': f'Request timed out after {self.timeout} seconds',
                'duration': time.time() - start_time
            }
        except requests.exceptions.ConnectionError:
            return {
                'summary': None,
                'summary_word_count': 0,
                'original_word_count': original_word_count,
                'success': False,
                'error': f'Cannot connect to Ollama server at {self.base_url}',
                'duration': time.time() - start_time
            }
        except requests.exceptions.HTTPError as e:
            return {
                'summary': None,
                'summary_word_count': 0,
                'original_word_count': original_word_count,
                'success': False,
                'error': f'HTTP error: {e.response.status_code} - {e.response.text[:100]}',
                'duration': time.time() - start_time
            }
        except Exception as e:
            return {
                'summary': None,
                'summary_word_count': 0,
                'original_word_count': original_word_count,
                'success': False,
                'error': f'Unexpected error: {str(e)}',
                'duration': time.time() - start_time
            }
    
    def translate_title(self, title, target_language='English'):
        """
        Translate article title to target language
        
        Args:
            title: Original title (typically German)
            target_language: Target language (default: 'English')
            
        Returns:
            {
                'success': bool,             # Whether translation succeeded
                'translated_title': str or None,  # Translated title
                'error': str or None,        # Error message if failed
                'duration': float            # Time taken in seconds
            }
        """
        if not self.enabled:
            return {
                'success': False,
                'translated_title': None,
                'error': 'Ollama is not enabled',
                'duration': 0
            }
        
        if not title or len(title.strip()) == 0:
            return {
                'success': False,
                'translated_title': None,
                'error': 'Title is empty',
                'duration': 0
            }
        
        start_time = time.time()
        
        try:
            # Construct messages for chat API
            messages = [
                {
                    'role': 'system',
                    'content': f"You are a professional translator. Translate the following German news headline to {target_language}.\\n\\nIMPORTANT: Provide ONLY the {target_language} translation. Do not include explanations, quotes, or any other text. Just the translated headline."
                },
                {
                    'role': 'user',
                    'content': title
                }
            ]
            
            # Make request using chat endpoint
            translated_title = self._chat_request(
                messages,
                options={
                    'temperature': 0.1,  # Low temperature for consistent translations
                    'num_predict': 100   # Limit response length
                }
            )
            
            if not translated_title:
                return {
                    'success': False,
                    'translated_title': None,
                    'error': 'Ollama returned empty translation',
                    'duration': time.time() - start_time
                }
            
            # Clean the translation output
            translated_title = self._clean_translation(translated_title)
            
            # Validate translation (if it's same as original, it might have failed)
            if translated_title.lower() == title.lower() and target_language == 'English':
                # Retry with more forceful prompt
                messages[0]['content'] += " If the text is already English, just output it as is."
                translated_title = self._chat_request(messages, options={'temperature': 0.1})
                translated_title = self._clean_translation(translated_title)
            
            return {
                'success': True,
                'translated_title': translated_title,
                'error': None,
                'duration': time.time() - start_time
            }
            
        except requests.exceptions.Timeout:
            return {
                'success': False,
                'translated_title': None,
                'error': f'Request timed out after {self.timeout} seconds',
                'duration': time.time() - start_time
            }
        except requests.exceptions.ConnectionError:
            return {
                'success': False,
                'translated_title': None,
                'error': f'Cannot connect to Ollama server at {self.base_url}',
                'duration': time.time() - start_time
            }
        except requests.exceptions.HTTPError as e:
            return {
                'success': False,
                'translated_title': None,
                'error': f'HTTP error: {e.response.status_code} - {e.response.text[:100]}',
                'duration': time.time() - start_time
            }
        except Exception as e:
            return {
                'success': False,
                'translated_title': None,
                'error': f'Unexpected error: {str(e)}',
                'duration': time.time() - start_time
            }
    
    def _clean_translation(self, translation):
        """Clean translation output by removing quotes and extra text"""
        # Extract first line only
        translation = translation.split('\n')[0]
        
        # Remove surrounding quotes (single and double)
        translation = translation.strip()
        if (translation.startswith('"') and translation.endswith('"')) or \
           (translation.startswith("'") and translation.endswith("'")):
            translation = translation[1:-1]
        
        # Trim whitespace again after quote removal
        translation = translation.strip()
        
        return translation
    
    def _clean_markdown(self, text):
        """Remove markdown formatting from text"""
        import re
        
        # Remove markdown headers (##, ###, etc.)
        text = re.sub(r'^#{1,6}\s+', '', text, flags=re.MULTILINE)
        
        # Remove bold/italic markers (**text**, *text*, __text__, _text_)
        text = re.sub(r'\*\*([^\*]+)\*\*', r'\1', text)
        text = re.sub(r'__([^_]+)__', r'\1', text)
        text = re.sub(r'\*([^\*]+)\*', r'\1', text)
        text = re.sub(r'_([^_]+)_', r'\1', text)
        
        # Remove markdown links [text](url) -> text
        text = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', text)
        
        # Remove inline code `text`
        text = re.sub(r'`([^`]+)`', r'\1', text)
        
        # Remove bullet points and list markers
        text = re.sub(r'^\s*[-*+]\s+', '', text, flags=re.MULTILINE)
        text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE)
        
        # Clean up extra whitespace
        text = re.sub(r'\n\s*\n', '\n\n', text)
        text = text.strip()
        
        return text
    
    def is_available(self):
        """
        Check if Ollama server is reachable
        
        Returns:
            bool: True if server is reachable, False otherwise
        """
        if not self.enabled:
            return False
        
        try:
            url = f"{self.base_url}/api/tags"
            headers = {}
            if self.api_key:
                headers['Authorization'] = f'Bearer {self.api_key}'
            
            response = requests.get(url, headers=headers, timeout=5)
            response.raise_for_status()
            return True
        except:
            return False
    
    def test_connection(self):
        """
        Test connection and return server info
        
        Returns:
            {
                'available': bool,
                'models': list,
                'current_model': str,
                'error': str or None
            }
        """
        if not self.enabled:
            return {
                'available': False,
                'models': [],
                'current_model': self.model,
                'error': 'Ollama is not enabled'
            }
        
        try:
            url = f"{self.base_url}/api/tags"
            headers = {}
            if self.api_key:
                headers['Authorization'] = f'Bearer {self.api_key}'
            
            response = requests.get(url, headers=headers, timeout=5)
            response.raise_for_status()
            
            result = response.json()
            models = [m.get('name', '') for m in result.get('models', [])]
            
            return {
                'available': True,
                'models': models,
                'current_model': self.model,
                'error': None
            }
        except requests.exceptions.ConnectionError:
            return {
                'available': False,
                'models': [],
                'current_model': self.model,
                'error': f'Cannot connect to Ollama server at {self.base_url}'
            }
        except Exception as e:
            return {
                'available': False,
                'models': [],
                'current_model': self.model,
                'error': str(e)
            }
    
    def generate(self, prompt, max_tokens=100):
        """
        Generate text using Ollama
        
        Args:
            prompt: Text prompt
            max_tokens: Maximum tokens to generate
            
        Returns:
            {
                'text': str,           # Generated text
                'success': bool,       # Whether generation succeeded
                'error': str or None,  # Error message if failed
                'duration': float      # Time taken in seconds
            }
        """
        if not self.enabled:
            return {
                'text': '',
                'success': False,
                'error': 'Ollama is disabled',
                'duration': 0
            }
        
        start_time = time.time()
        
        try:
            messages = [{'role': 'user', 'content': prompt}]
            
            text = self._chat_request(
                messages,
                options={
                    "num_predict": max_tokens,
                    "temperature": 0.1
                }
            )
            
            duration = time.time() - start_time
            
            return {
                'text': text,
                'success': True,
                'error': None,
                'duration': duration
            }
                
        except requests.exceptions.Timeout:
            return {
                'text': '',
                'success': False,
                'error': f"Request timed out after {self.timeout}s",
                'duration': time.time() - start_time
            }
        except Exception as e:
            return {
                'text': '',
                'success': False,
                'error': str(e),
                'duration': time.time() - start_time
            }
    
    def extract_keywords(self, title, summary, max_keywords=5):
        """
        Extract keywords/topics from article for personalization
        
        Args:
            title: Article title
            summary: Article summary
            max_keywords: Maximum number of keywords to extract (default 5)
            
        Returns:
            {
                'keywords': list,      # List of extracted keywords
                'success': bool,       # Whether extraction succeeded
                'error': str or None,  # Error message if failed
                'duration': float      # Time taken in seconds
            }
        """
        if not self.enabled:
            return {
                'keywords': [],
                'success': False,
                'error': 'Ollama is disabled',
                'duration': 0
            }
        
        start_time = time.time()
        
        try:
            # Construct messages for chat API
            messages = [
                {
                    'role': 'system',
                    'content': f"Extract {max_keywords} key topics or keywords from the article.\\n\\nReturn ONLY the keywords separated by commas, nothing else. Focus on:\\n- Main topics (e.g., 'Bayern Munich', 'Oktoberfest', 'City Council')\\n- Locations (e.g., 'Marienplatz', 'Airport')\\n- Events or themes (e.g., 'Transportation', 'Housing', 'Technology')"
                },
                {
                    'role': 'user',
                    'content': f"Title: {title}\\nSummary: {summary}"
                }
            ]
            
            # Make request
            keywords_text = self._chat_request(
                messages,
                options={
                    'temperature': 0.2,
                    'num_predict': 100
                }
            )
            
            if not keywords_text:
                return {
                    'keywords': [],
                    'success': False,
                    'error': 'Ollama returned empty response',
                    'duration': time.time() - start_time
                }
            
            # Parse keywords from response
            keywords = [k.strip() for k in keywords_text.split(',')]
            keywords = [k for k in keywords if k and len(k) > 2][:max_keywords]
            
            return {
                'keywords': keywords,
                'success': True,
                'error': None,
                'duration': time.time() - start_time
            }
            
        except requests.exceptions.Timeout:
            return {
                'keywords': [],
                'success': False,
                'error': f"Request timed out after {self.timeout}s",
                'duration': time.time() - start_time
            }
        except Exception as e:
            return {
                'keywords': [],
                'success': False,
                'error': str(e),
                'duration': time.time() - start_time
            }


if __name__ == '__main__':
    # Quick test
    import os
    from dotenv import load_dotenv
    
    load_dotenv(dotenv_path='../.env')
    
    client = OllamaClient(
        base_url=os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434'),
        model=os.getenv('OLLAMA_MODEL', 'phi3:latest'),
        enabled=True
    )
    
    print("Testing Ollama connection...")
    result = client.test_connection()
    print(f"Available: {result['available']}")
    print(f"Models: {result['models']}")
    print(f"Current model: {result['current_model']}")
    
    if result['available']:
        print("\nTesting summarization...")
        test_content = """
        The new U-Bahn line connecting Munich's city center with the airport opened today.
        Mayor Dieter Reiter attended the opening ceremony along with hundreds of residents.
        The line will significantly reduce travel time between the airport and downtown Munich.
        Construction took five years and cost approximately 2 billion euros.
        The new line includes 10 stations and runs every 10 minutes during peak hours.
        """
        
        summary_result = client.summarize_article(test_content, max_words=50)
        print(f"Success: {summary_result['success']}")
        print(f"Summary: {summary_result['summary']}")
        print(f"Original word count: {summary_result['original_word_count']}")
        print(f"Summary word count: {summary_result['summary_word_count']}")
        print(f"Compression: {summary_result['original_word_count'] / max(summary_result['summary_word_count'], 1):.1f}x")
        print(f"Duration: {summary_result['duration']:.2f}s")