"""
AI processor module for FactsDB service
Handles communication with OpenAI compatible endpoints
"""

import requests
import json
from typing import Dict, Any, Optional
from .config import AIEndpointConfig

class AIEndpointClient:
    """Client for communicating with OpenAI compatible endpoint"""
    
    def __init__(self, config: AIEndpointConfig):
        self.config = config
        self.base_url = config.url.rstrip('/')
        self.auth_token = config.auth_token
    
    def _get_headers(self) -> Dict[str, str]:
        """Get headers with authentication"""
        return {
            'Authorization': f'Bearer {self.auth_token}',
            'Content-Type': 'application/json'
        }
    
    def send_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Send request to AI endpoint"""
        url = f"{self.base_url}/v1/chat/completions"
        
        try:
            response = requests.post(
                url,
                headers=self._get_headers(),
                json=payload,
                timeout=300  # 5 minute timeout
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise Exception(f"AI endpoint request failed: {str(e)}")
    
    def extract_facts(self, text_content: str, prompt: str, model: str = "gpt-oss") -> Dict[str, Any]:
        """Extract facts from text using AI with chunking support"""
        chunk_size = 8000
        overlap = 500

        # Default prompt
        default_prompt = """Extract key facts from the following article in structured JSON format.
Return only valid JSON without any additional text.

Extract the following information:
1. Key entities (companies, people, locations, organizations)
2. Key dates or time periods mentioned
3. Main facts from the article

Format the response as a JSON object with these fields:
{
    "fact": "main fact extracted from the article",
    "key_entities": ["entity1", "entity2"],
    "key_dates": ["date1", "date2"]
}"""

        final_prompt = prompt if prompt else default_prompt

        # Split long content into overlapping chunks
        chunks = []
        if len(text_content) > chunk_size:
            start = 0
            while start < len(text_content):
                end = min(start + chunk_size, len(text_content))
                chunks.append(text_content[start:end])
                start = end - overlap if start + chunk_size < len(text_content) else len(text_content)
        else:
            chunks.append(text_content)

        all_entities = set()
        all_dates = set()
        all_facts = []

        for i, chunk in enumerate(chunks):
            payload = {
                "model": model,
                "messages": [
                    {
                        "role": "user",
                        "content": f"{final_prompt}\n\nArticle Content (part {i + 1}/{len(chunks)}):\n{chunk}"
                    }
                ],
                "temperature": 0.3,
                "max_tokens": 1000
            }

            try:
                response = self.send_request(payload)
                if 'choices' in response and len(response['choices']) > 0:
                    response_text = response['choices'][0]['message']['content'].strip()
                    if response_text.startswith('```json'):
                        response_text = response_text[7:-3].strip()
                    elif response_text.startswith('```'):
                        response_text = response_text[3:-3].strip()
                    result = json.loads(response_text)
                    if isinstance(result.get('key_entities'), list):
                        all_entities.update(result['key_entities'])
                    if isinstance(result.get('key_dates'), list):
                        all_dates.update(result['key_dates'])
                    if result.get('fact'):
                        all_facts.append(result['fact'])
            except Exception as e:
                print(f"Warning: Failed to process chunk {i + 1}: {str(e)}")
                continue

        return {
            "fact": ". ".join(all_facts) if all_facts else "No facts extracted",
            "key_entities": list(all_entities),
            "key_dates": list(all_dates)
        }


class AIProcessor:
    """Main AI processor class for FactsDB"""
    
    def __init__(self, config: AIEndpointConfig):
        self.client = AIEndpointClient(config)
    
    def extract_facts_from_text(self, text_content: str, prompt: str = "", model: str = "gpt-oss") -> Dict[str, Any]:
        """Extract facts from text content using AI"""
        return self.client.extract_facts(text_content, prompt, model)
    
    def validate_model_support(self, model: str) -> bool:
        """Validate if model is supported"""
        # In a real implementation, this would check against available models
        supported_models = ["gpt-oss", "qwen3", "qwen3-coder"]
        return model in supported_models