"""
Ollama Client for gpt-oss:20b Integration
Local LLM inference for ADAM-1 multi-agent system

Model: gpt-oss:20b
- 20.9B parameters
- 128K context window
- Chain-of-thought reasoning
- Apache 2.0 license
- Requires ≥16GB VRAM
"""

import json
import logging
from dataclasses import dataclass
from typing import Dict, Generator, List, Optional, Union

import requests

logger = logging.getLogger(__name__)


@dataclass
class OllamaResponse:
    """Response from Ollama API."""
    content: str
    model: str
    done: bool
    total_duration: Optional[int] = None
    prompt_eval_count: Optional[int] = None
    eval_count: Optional[int] = None


class OllamaClient:
    """
    Client for interacting with Ollama API.

    Features:
    - Synchronous and streaming inference
    - Chain-of-thought prompting
    - Context window management (128K tokens)
    - Automatic retry logic
    - Token usage tracking
    """

    def __init__(
        self,
        model: str = "gpt-oss:20b",
        base_url: str = "http://localhost:11434",
        timeout: int = 300,
        max_retries: int = 3,
        num_ctx: int = 131072
    ):
        """
        Initialize Ollama client.

        Args:
            model: Model name (default: gpt-oss:20b)
            base_url: Ollama API base URL
            timeout: Request timeout in seconds
            max_retries: Maximum retry attempts
            num_ctx: Context window size in tokens (default: 131072 for 128K)
                     CRITICAL: Without this, Ollama defaults to 8192 tokens
                     and silently truncates long prompts!
        """
        self.model = model
        self.base_url = base_url
        self.timeout = timeout
        self.max_retries = max_retries
        self.num_ctx = num_ctx

        logger.info(f"Initialized Ollama client: {model} @ {base_url} (num_ctx={num_ctx})")

    def is_available(self) -> bool:
        """
        Check if Ollama service is available.

        Returns:
            True if service is running
        """
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            return response.status_code == 200
        except Exception as e:
            logger.warning(f"Ollama service not available: {e}")
            return False

    def list_models(self) -> List[str]:
        """
        List available models.

        Returns:
            List of model names
        """
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            response.raise_for_status()
            data = response.json()
            return [model["name"] for model in data.get("models", [])]
        except Exception as e:
            logger.error(f"Failed to list models: {e}")
            return []

    def generate(
        self,
        prompt: str,
        system: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        stream: bool = False
    ) -> Union[OllamaResponse, Generator[str, None, None]]:
        """
        Generate completion from prompt.

        Args:
            prompt: User prompt
            system: System prompt (optional)
            temperature: Sampling temperature (0.0-1.0)
            max_tokens: Maximum tokens to generate
            stream: Stream response

        Returns:
            OllamaResponse or generator of response chunks
        """
        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": stream,
            "options": {
                "temperature": temperature,
                "num_ctx": self.num_ctx,
            }
        }

        if system:
            payload["system"] = system

        if max_tokens:
            payload["options"]["num_predict"] = max_tokens

        logger.debug(f"Generating with {self.model}: {prompt[:100]}...")

        try:
            response = requests.post(
                f"{self.base_url}/api/generate",
                json=payload,
                timeout=self.timeout,
                stream=stream
            )
            response.raise_for_status()

            if stream:
                return self._stream_response(response)
            else:
                data = response.json()
                return OllamaResponse(
                    content=data["response"],
                    model=data["model"],
                    done=data["done"],
                    total_duration=data.get("total_duration"),
                    prompt_eval_count=data.get("prompt_eval_count"),
                    eval_count=data.get("eval_count")
                )

        except Exception as e:
            logger.error(f"Generation failed: {e}")
            raise

    def _stream_response(self, response: requests.Response) -> Generator[str, None, None]:
        """
        Stream response chunks.

        Args:
            response: Streaming HTTP response

        Yields:
            Response text chunks
        """
        for line in response.iter_lines():
            if line:
                data = json.loads(line)
                if "response" in data:
                    yield data["response"]
                if data.get("done", False):
                    break

    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        stream: bool = False
    ) -> Union[OllamaResponse, Generator[str, None, None]]:
        """
        Chat completion with message history.

        Args:
            messages: List of message dicts with 'role' and 'content'
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            stream: Stream response

        Returns:
            OllamaResponse or generator of response chunks
        """
        payload = {
            "model": self.model,
            "messages": messages,
            "stream": stream,
            "options": {
                "temperature": temperature,
                "num_ctx": self.num_ctx,
            }
        }

        if max_tokens:
            payload["options"]["num_predict"] = max_tokens

        logger.debug(f"Chat with {len(messages)} messages")

        try:
            response = requests.post(
                f"{self.base_url}/api/chat",
                json=payload,
                timeout=self.timeout,
                stream=stream
            )
            response.raise_for_status()

            if stream:
                return self._stream_chat_response(response)
            else:
                data = response.json()
                return OllamaResponse(
                    content=data["message"]["content"],
                    model=data["model"],
                    done=data["done"],
                    total_duration=data.get("total_duration"),
                    prompt_eval_count=data.get("prompt_eval_count"),
                    eval_count=data.get("eval_count")
                )

        except Exception as e:
            logger.error(f"Chat failed: {e}")
            raise

    def _stream_chat_response(self, response: requests.Response) -> Generator[str, None, None]:
        """
        Stream chat response chunks.

        Args:
            response: Streaming HTTP response

        Yields:
            Response text chunks
        """
        for line in response.iter_lines():
            if line:
                data = json.loads(line)
                if "message" in data and "content" in data["message"]:
                    yield data["message"]["content"]
                if data.get("done", False):
                    break

    def adam_summarization_query(
        self,
        context: str,
        clinical_features: str,
        microbiome_features: str,
        temperature: float = 0.3
    ) -> str:
        """
        ADAM-1 summarization agent query.

        Args:
            context: Retrieved literature context
            clinical_features: Clinical data summary
            microbiome_features: Microbiome data summary
            temperature: Lower temperature for factual summarization

        Returns:
            Summarization response
        """
        system_prompt = (
            "You are a medical research summarization agent specializing in Alzheimer's disease.\n"
            "Your task is to synthesize clinical data, microbiome data, and scientific literature\n"
            "to provide comprehensive summaries for diagnostic classification.\n\n"
            "Focus on:\n"
            "- Alzheimer's disease biomarkers (amyloid, tau, APOE)\n"
            "- Gut microbiome composition and metabolites\n"
            "- Clinical correlations and mechanisms\n"
            "- Evidence-based conclusions"
        )

        user_prompt = f"""Summarize the following data for Alzheimer's disease classification:

CLINICAL FEATURES:
{clinical_features}

MICROBIOME FEATURES:
{microbiome_features}

SCIENTIFIC LITERATURE CONTEXT:
{context}

Provide a comprehensive summary highlighting key findings, correlations, \
and diagnostic implications."""

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        logger.info("ADAM summarization query")
        response = self.chat(messages, temperature=temperature)
        return response.content

    def adam_classification_query(
        self,
        summary: str,
        context: str,
        temperature: float = 0.1
    ) -> Dict[str, any]:
        """
        ADAM-1 classification agent query.

        Args:
            summary: Output from summarization agent
            context: Additional literature context
            temperature: Very low temperature for consistent classification

        Returns:
            Classification result with confidence
        """
        system_prompt = (
            "You are a medical classification agent specializing in "
            "Alzheimer's disease diagnosis.\n"
            "Based on summarized clinical and microbiome data, along with scientific literature,\n"
            "classify the patient's Alzheimer's disease status.\n\n"
            "Classification categories:\n"
            "- Healthy Control (HC)\n"
            "- Mild Cognitive Impairment (MCI)\n"
            "- Alzheimer's Disease (AD)\n\n"
            "Provide:\n"
            "1. Classification label\n"
            "2. Confidence score (0-1)\n"
            "3. Key supporting evidence\n"
            "4. Diagnostic reasoning"
        )

        user_prompt = f"""Classify the following patient based on the summary and literature:

SUMMARY:
{summary}

ADDITIONAL LITERATURE:
{context}

Provide classification in JSON format:
{{
  "classification": "HC/MCI/AD",
  "confidence": 0.0-1.0,
  "evidence": ["key point 1", "key point 2", ...],
  "reasoning": "diagnostic reasoning"
}}"""

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        logger.info("ADAM classification query")
        response = self.chat(messages, temperature=temperature)

        # Try to parse JSON response
        try:
            result = json.loads(response.content)
            return result
        except json.JSONDecodeError:
            logger.warning("Failed to parse JSON response, returning raw text")
            return {
                "classification": "UNKNOWN",
                "confidence": 0.0,
                "evidence": [],
                "reasoning": response.content
            }


def create_ollama_client(model: str = "gpt-oss:20b") -> OllamaClient:
    """
    Factory function to create Ollama client.

    Args:
        model: Model name

    Returns:
        OllamaClient instance
    """
    return OllamaClient(model=model)


if __name__ == "__main__":
    # Test the client
    logging.basicConfig(level=logging.INFO)

    print("=" * 60)
    print("Ollama Client Test")
    print("=" * 60)

    client = OllamaClient()

    # Check availability
    print("\n1. Checking Ollama service...")
    if client.is_available():
        print("   ✅ Ollama service is running")

        # List models
        print("\n2. Available models:")
        models = client.list_models()
        for model in models:
            print(f"   - {model}")

        # Test generation (only if gpt-oss:20b is available)
        if "gpt-oss:20b" in models:
            print("\n3. Testing generation...")
            response = client.generate(
                prompt="What are the main biomarkers of Alzheimer's disease?",
                temperature=0.3
            )
            print(f"   Response: {response.content[:200]}...")
            print(f"   Tokens: {response.eval_count}")
        else:
            print("\n3. ⚠️  gpt-oss:20b not found")
            print("   Run: ollama pull gpt-oss:20b")
    else:
        print("   ❌ Ollama service not running")
        print("   Please install Ollama and run: sudo bash INSTALL_OLLAMA.sh")

    print("\n" + "=" * 60)
