"""
Ollama integration service for AI-based speaker detection and text segmentation.
"""
import json
import logging
import os
from typing import List, Optional, Dict, Tuple
import httpx


logger = logging.getLogger(__name__)


class OllamaService:
    """
    Service for integrating with Ollama API for AI-based speaker detection.
    
    Uses llama3:8b model to intelligently detect speakers in text without
    explicit role markers.
    """
    
    def __init__(self, base_url: Optional[str] = None):
        """
        Initializes Ollama service with base URL.
        
        Args:
            base_url: Ollama API base URL. Defaults to OLLAMA_BASE_URL env var
                     or http://localhost:11434 if not set.
        """
        self.base_url = base_url or os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        self.model = "llama3:latest"  # Using llama3:latest (8B model)
        self.timeout = 30.0  # 30 seconds timeout
        
    async def detect_speakers(self, text: str) -> List[str]:
        """
        Sends text to Ollama for speaker detection.
        Samples up to 4 chunks of 5000 characters from different parts of the text
        to ensure comprehensive speaker detection in long texts.
        
        Args:
            text: Input text to analyze for speakers
            
        Returns:
            List of detected speaker names (deduplicated across all samples)
            
        Raises:
            Exception: If Ollama is unavailable or request fails
        """
        # Get text samples from different parts of the text
        samples = self._sample_text(text, chunk_size=5000, num_samples=4)
        
        all_speakers = set()
        
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                # Process each sample
                for i, sample in enumerate(samples):
                    prompt = self._build_prompt(sample)
                    
                    logger.debug(f"Sending sample {i+1}/{len(samples)} to Ollama ({len(sample)} chars)")
                    
                    response = await client.post(
                        f"{self.base_url}/api/generate",
                        json={
                            "model": self.model,
                            "prompt": prompt,
                            "stream": False,
                            "format": "json"
                        }
                    )
                    response.raise_for_status()
                    
                    result = response.json()
                    response_text = result.get("response", "")
                    
                    speakers = self._parse_response(response_text)
                    logger.info(f"Sample {i+1} detected {len(speakers)} speakers: {speakers}")
                    
                    # Add to our set (automatically deduplicates)
                    all_speakers.update(speakers)
                
                speaker_list = sorted(list(all_speakers))
                logger.info(f"Ollama detected {len(speaker_list)} total unique speakers: {speaker_list}")
                return speaker_list
                
        except httpx.TimeoutException:
            logger.warning("Ollama request timed out after 30 seconds")
            raise Exception("Ollama request timed out")
        except httpx.HTTPError as e:
            logger.warning(f"Ollama HTTP error: {e}")
            raise Exception(f"Ollama service error: {e}")
        except Exception as e:
            logger.warning(f"Ollama detection failed: {e}")
            raise
    
    def _sample_text(self, text: str, chunk_size: int = 5000, num_samples: int = 4) -> List[str]:
        """
        Samples text from different parts to ensure comprehensive speaker detection.
        
        For short texts, returns the full text.
        For long texts, samples from beginning, middle sections, and end.
        
        Args:
            text: Input text to sample
            chunk_size: Size of each sample in characters
            num_samples: Number of samples to extract
            
        Returns:
            List of text samples
        """
        text_len = len(text)
        
        # If text is shorter than one chunk, return it as-is
        if text_len <= chunk_size:
            return [text]
        
        # If text is shorter than what we'd sample, return fewer samples
        max_possible_samples = (text_len // chunk_size) + 1
        actual_samples = min(num_samples, max_possible_samples)
        
        samples = []
        
        if actual_samples == 1:
            # Just take the beginning
            samples.append(text[:chunk_size])
        else:
            # Distribute samples evenly across the text
            step = text_len // actual_samples
            
            for i in range(actual_samples):
                start = i * step
                end = min(start + chunk_size, text_len)
                samples.append(text[start:end])
        
        return samples
    
    def _build_prompt(self, text: str) -> str:
        """
        Constructs prompt for speaker detection.
        
        Instructs model to return JSON array of speaker names.
        
        Args:
            text: Input text to analyze (should already be sampled/truncated)
            
        Returns:
            Formatted prompt string
        """
        prompt = f"""Analyze the following text and identify all distinct speakers or characters who need voice assignments.

CRITICAL: Look for TWO types of content:
1. NARRATIVE TEXT: Descriptions, scene-setting, "he said", "she thought", etc. → This needs "Narrator"
2. CHARACTER DIALOGUE: Direct speech by characters → These need character names

You MUST include "Narrator" if there is ANY narrative text, descriptions, or non-dialogue content.

Return ONLY a JSON object with a "speakers" array.

Examples:
- Text: "Alice said hello. Bob waved back." → {{"speakers": ["Narrator", "Alice", "Bob"]}}
- Text: "Once upon a time, there was a princess." → {{"speakers": ["Narrator"]}}
- Text: "The wizard spoke. 'I see danger,' he warned." → {{"speakers": ["Narrator"]}}
- Text: "Alice: Hello! Bob: Hi there!" → {{"speakers": ["Alice", "Bob"]}}

Text to analyze:
{text}

Return format: {{"speakers": ["Narrator", "Character1", "Character2", ...]}}
"""
        return prompt
    
    def _parse_response(self, response: str) -> List[str]:
        """
        Parses Ollama response to extract speaker names.
        
        Handles both JSON and plain text responses.
        
        Args:
            response: Response text from Ollama
            
        Returns:
            List of speaker names
        """
        if not response or not response.strip():
            return []
        
        try:
            # Try to parse as JSON
            data = json.loads(response.strip())
            
            # Handle different possible JSON structures
            if isinstance(data, dict):
                if "speakers" in data:
                    speakers = data["speakers"]
                elif "characters" in data:
                    speakers = data["characters"]
                elif "roles" in data:
                    speakers = data["roles"]
                else:
                    # Try to find any array in the response
                    for value in data.values():
                        if isinstance(value, list):
                            speakers = value
                            break
                    else:
                        return []
            elif isinstance(data, list):
                speakers = data
            else:
                return []
            
            # Ensure we have a list of strings
            if isinstance(speakers, list):
                return [str(s).strip() for s in speakers if s]
            
            return []
            
        except json.JSONDecodeError:
            logger.warning(f"Failed to parse Ollama response as JSON: {response[:100]}")
            # Try to extract speaker names from plain text
            # Look for common patterns like lists or quoted names
            return []
    
    async def segment_text(self, text: str, speakers: List[str]) -> List[Tuple[str, str]]:
        """
        Uses AI to segment text and assign each segment to a speaker.
        
        Args:
            text: Input text to segment
            speakers: List of known speakers in the text
            
        Returns:
            List of tuples: (speaker_name, text_segment)
            
        Raises:
            Exception: If Ollama is unavailable or request fails
        """
        try:
            async with httpx.AsyncClient(timeout=60.0) as client:  # Longer timeout for segmentation
                prompt = self._build_segmentation_prompt(text, speakers)
                
                logger.info(f"Requesting AI segmentation for {len(text)} chars with {len(speakers)} speakers")
                
                response = await client.post(
                    f"{self.base_url}/api/generate",
                    json={
                        "model": self.model,
                        "prompt": prompt,
                        "stream": False,
                        "format": "json"
                    }
                )
                response.raise_for_status()
                
                result = response.json()
                response_text = result.get("response", "")
                
                segments = self._parse_segmentation_response(response_text)
                logger.info(f"AI segmentation produced {len(segments)} segments")
                
                return segments
                
        except httpx.TimeoutException:
            logger.warning("Ollama segmentation request timed out")
            raise Exception("Ollama segmentation timed out")
        except httpx.HTTPError as e:
            logger.warning(f"Ollama HTTP error during segmentation: {e}")
            raise Exception(f"Ollama service error: {e}")
        except Exception as e:
            logger.warning(f"Ollama segmentation failed: {e}")
            raise
    
    def _build_segmentation_prompt(self, text: str, speakers: List[str]) -> str:
        """
        Constructs prompt for text segmentation with speaker attribution.
        
        Args:
            text: Input text to segment
            speakers: List of known speakers
            
        Returns:
            Formatted prompt string
        """
        speakers_str = ", ".join(f'"{s}"' for s in speakers)
        
        prompt = f"""You are analyzing a story or dialogue. Your task is to split the text into segments and assign each segment to the appropriate speaker.

Known speakers: {speakers_str}

Rules:
1. Split the text into logical segments (sentences or short paragraphs)
2. Assign each segment to the speaker who is speaking or being described
3. Use "Narrator" for narrative text, descriptions, or text that doesn't belong to a specific character
4. If a character is mentioned but not speaking, use "Narrator"
5. Only use speakers from the known speakers list

Return a JSON object with a "segments" array. Each segment should have:
- "speaker": the name of the speaker (must be from the known speakers list or "Narrator")
- "text": the text content for that segment

Example format:
{{
  "segments": [
    {{"speaker": "Narrator", "text": "The old wizard gazed into his crystal ball."}},
    {{"speaker": "Wizard", "text": "I see danger approaching."}},
    {{"speaker": "Narrator", "text": "A young knight heard the warning."}}
  ]
}}

Text to segment:
{text}

Return ONLY the JSON object with the segments array.
"""
        return prompt
    
    def _parse_segmentation_response(self, response: str) -> List[Tuple[str, str]]:
        """
        Parses Ollama segmentation response to extract speaker-text pairs.
        
        Args:
            response: Response text from Ollama
            
        Returns:
            List of tuples: (speaker_name, text_content)
        """
        if not response or not response.strip():
            return []
        
        try:
            # Try to parse as JSON
            data = json.loads(response.strip())
            
            # Extract segments array
            segments_data = []
            if isinstance(data, dict):
                if "segments" in data:
                    segments_data = data["segments"]
                else:
                    # Try to find any array in the response
                    for value in data.values():
                        if isinstance(value, list):
                            segments_data = value
                            break
            elif isinstance(data, list):
                segments_data = data
            
            # Convert to tuples
            segments = []
            for seg in segments_data:
                if isinstance(seg, dict) and "speaker" in seg and "text" in seg:
                    speaker = str(seg["speaker"]).strip()
                    text = str(seg["text"]).strip()
                    if speaker and text:
                        segments.append((speaker, text))
            
            return segments
            
        except json.JSONDecodeError:
            logger.warning(f"Failed to parse Ollama segmentation response as JSON: {response[:200]}")
            return []
    
    async def is_available(self) -> bool:
        """
        Checks if Ollama service is available.
        
        Returns:
            True if service responds to health check, False otherwise
        """
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                response = await client.get(f"{self.base_url}/api/tags")
                return response.status_code == 200
        except Exception as e:
            logger.debug(f"Ollama availability check failed: {e}")
            return False