"""Base agent class."""
import asyncio
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlsplit, urlunsplit
from anthropic import AnthropicFoundry
from openai import OpenAI
import httpx
from loguru import logger

from src.config import settings, estimate_cost
from src.database.client import db


class BaseAgent(ABC):
    """Base class for all AI agents."""
    @staticmethod
    def _normalize_azure_claude_base_url(base_url: str) -> str:
        """Normalize Azure Claude URL to the AnthropicFoundry base URL."""
        raw = (base_url or "").strip()
        if not raw:
            return raw

        parsed = urlsplit(raw)
        scheme = parsed.scheme or "https"
        netloc = parsed.netloc or parsed.path
        path = parsed.path if parsed.netloc else ""
        path = (path or "").rstrip("/")

        if path.endswith("/anthropic/v1/messages"):
            normalized_path = path[: -len("/v1/messages")]
        elif path.endswith("/anthropic"):
            normalized_path = path
        else:
            for suffix in ("/openai/v1/chat/completions", "/openai/v1", "/openai", "/models"):
                if path.endswith(suffix):
                    path = path[: -len(suffix)]
                    break
            normalized_path = f"{path}/anthropic"

        return urlunsplit((scheme, netloc, normalized_path, "", ""))

    def __init__(self, name: str):
        """
        Initialize base agent.

        Args:
            name: Name of the agent
        """
        self.name = name
        self.openai_client = OpenAI(api_key=settings.openai_api_key)
        self.azure_claude_client: Optional[Any] = None
        self.azure_claude_base_url: str = ""
        if settings.azure_claude_api_key:
            base_source = settings.azure_claude_base_url or settings.azure_claude_endpoint
            if base_source:
                self.azure_claude_base_url = self._normalize_azure_claude_base_url(base_source)
                self.azure_claude_client = AnthropicFoundry(
                    api_key=settings.azure_claude_api_key,
                    base_url=self.azure_claude_base_url
                )
                logger.info(f"[{self.name}] Azure Claude base URL normalized to {self.azure_claude_base_url}")
        self._usage_logs: List[Dict[str, Any]] = []
        self._user_id: Optional[str] = None
        self._company_id: Optional[str] = None
        self._operation: str = "unknown"
        logger.info(f"Initialized {name} agent")

    def set_tracking_context(
        self,
        operation: str,
        user_id: Optional[str] = None,
        company_id: Optional[str] = None
    ):
        """Set context for usage tracking."""
        self._operation = operation
        self._user_id = user_id
        self._company_id = company_id
        self._usage_logs = []

    @abstractmethod
    async def process(self, *args, **kwargs) -> Any:
        """Process the agent's task."""
        pass

    async def _log_usage(self, provider: str, model: str, prompt_tokens: int, completion_tokens: int, total_tokens: int):
        """Log API usage to database."""
        cost = estimate_cost(model, prompt_tokens, completion_tokens)
        usage = {
            "provider": provider,
            "model": model,
            "operation": self._operation,
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
            "estimated_cost_usd": cost
        }
        self._usage_logs.append(usage)

        try:
            from uuid import UUID
            await db.log_api_usage(
                provider=provider,
                model=model,
                operation=self._operation,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=total_tokens,
                estimated_cost_usd=cost,
                user_id=UUID(self._user_id) if self._user_id else None,
                company_id=UUID(self._company_id) if self._company_id else None
            )
        except Exception as e:
            logger.warning(f"Failed to log usage to DB: {e}")

        # Increment company token quota
        if self._company_id:
            try:
                from uuid import UUID
                await db.increment_company_tokens(UUID(self._company_id), total_tokens)
            except Exception as e:
                logger.warning(f"Failed to increment company tokens: {e}")

    async def call_openai(
        self,
        system_prompt: str,
        user_prompt: str,
        model: str = "gpt-4o",
        temperature: float = 0.7,
        response_format: Optional[Dict[str, str]] = None
    ) -> str:
        """
        Call OpenAI API.

        Args:
            system_prompt: System message
            user_prompt: User message
            model: Model to use
            temperature: Temperature for sampling
            response_format: Optional response format (e.g., {"type": "json_object"})

        Returns:
            Assistant's response
        """
        logger.info(f"[{self.name}] Calling OpenAI ({model})")

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

        kwargs = {
            "model": model,
            "messages": messages,
            "temperature": temperature
        }

        if response_format:
            kwargs["response_format"] = response_format

        # Run synchronous OpenAI call in thread pool to avoid blocking event loop
        response = await asyncio.to_thread(
            self.openai_client.chat.completions.create,
            **kwargs
        )

        result = response.choices[0].message.content
        logger.debug(f"[{self.name}] Received response (length: {len(result)})")

        # Track usage
        if response.usage:
            await self._log_usage(
                provider="openai",
                model=model,
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
                total_tokens=response.usage.total_tokens
            )

        return result

    async def call_azure_claude(
        self,
        system_prompt: str,
        user_prompt: str,
        model: Optional[str] = None,
        temperature: float = 0.7,
        response_format: Optional[Dict[str, str]] = None
    ) -> str:
        """
        Call Azure-hosted Claude via the Anthropic-compatible Azure endpoint.

        Args:
            system_prompt: System message
            user_prompt: User message
            model: Human-readable model identifier for logging/pricing
            temperature: Temperature for sampling
            response_format: Optional structured response format

        Returns:
            Assistant's response
        """
        if not self.azure_claude_client or not settings.azure_claude_api_key:
            raise ValueError(
                "Azure Claude is not configured. Set AZURE_CLAUDE_API_KEY and either "
                "AZURE_CLAUDE_BASE_URL or AZURE_CLAUDE_ENDPOINT."
            )

        display_model = model or settings.azure_claude_model
        deployment = settings.azure_claude_deployment or display_model
        logger.info(f"[{self.name}] Calling Azure Claude ({display_model} via {deployment})")

        kwargs: Dict[str, Any] = {
            "model": deployment,
            "system": system_prompt,
            "messages": [
                {"role": "user", "content": user_prompt}
            ],
            "temperature": temperature,
            "max_tokens": 4096
        }

        if response_format:
            kwargs["system"] = (
                system_prompt
                + "\n\nWICHTIG: Antworte mit einem einzigen gültigen JSON-Objekt und ohne zusätzliche Erklärung."
            )

        response = await asyncio.to_thread(
            self.azure_claude_client.messages.create,
            **kwargs
        )

        content_blocks = getattr(response, "content", []) or []
        text_parts = []
        for block in content_blocks:
            if hasattr(block, "text"):
                text_parts.append(block.text or "")
            elif isinstance(block, dict) and block.get("type") == "text":
                text_parts.append(block.get("text", ""))
        final_text = "\n".join(part for part in text_parts if part).strip()
        logger.debug(f"[{self.name}] Received Azure Claude response (length: {len(final_text)})")

        usage = getattr(response, "usage", None)
        input_tokens = getattr(usage, "input_tokens", 0) if usage else 0
        output_tokens = getattr(usage, "output_tokens", 0) if usage else 0
        total_tokens = input_tokens + output_tokens
        if total_tokens:
            await self._log_usage(
                provider="azure",
                model=display_model,
                prompt_tokens=input_tokens,
                completion_tokens=output_tokens,
                total_tokens=total_tokens
            )

        return final_text

    async def call_perplexity(
        self,
        system_prompt: str,
        user_prompt: str,
        model: str = "sonar"
    ) -> str:
        """
        Call Perplexity API for research.

        Args:
            system_prompt: System message
            user_prompt: User message
            model: Model to use

        Returns:
            Assistant's response
        """
        logger.info(f"[{self.name}] Calling Perplexity ({model})")

        url = "https://api.perplexity.ai/chat/completions"
        headers = {
            "Authorization": f"Bearer {settings.perplexity_api_key}",
            "Content-Type": "application/json"
        }

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

        async with httpx.AsyncClient() as client:
            response = await client.post(url, json=payload, headers=headers, timeout=60.0)
            response.raise_for_status()
            result = response.json()

        content = result["choices"][0]["message"]["content"]
        logger.debug(f"[{self.name}] Received Perplexity response (length: {len(content)})")

        # Track usage
        usage = result.get("usage", {})
        if usage:
            await self._log_usage(
                provider="perplexity",
                model=model,
                prompt_tokens=usage.get("prompt_tokens", 0),
                completion_tokens=usage.get("completion_tokens", 0),
                total_tokens=usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
            )

        return content
