import httpx
import logging
import asyncio
import datetime
from datetime import UTC
from typing import Any, Optional, Dict, List, Union
import weakref

from .config import (
    ClientSettings,
    AuthConfig, 
    AuthType, 
    BasicAuthConfig, 
    ApiKeyAuthConfig, 
    ApiKeyLocation,
    OAuth2AuthConfig
)

# Observability imports from beamflow_lib
from beamflow_lib.observability.model import DataExchangeEvent
from beamflow_lib.observability.ingestion import record_data_exchange

logger = logging.getLogger(__name__)

class OAuth2Manager:
    """
    Primitive OAuth2 manager for token handling and re-auth.
    """
    def __init__(self, config: OAuth2AuthConfig):
        self.config = config
        self._access_token = None
        self._lock = asyncio.Lock()

    async def get_token(self, client: httpx.AsyncClient) -> str:
        async with self._lock:
            if self._access_token:
                return self._access_token
            return await self.refresh_token(client)

    async def refresh_token(self, client: httpx.AsyncClient) -> str:
        logger.info(f"Refreshing OAuth2 token for {self.config.client_id} via {self.config.token_url}")
        data = {
            "grant_type": "client_credentials",
            "client_id": self.config.client_id,
            "client_secret": self.config.client_secret,
            "scope": " ".join(self.config.scopes)
        }
        if self.config.refresh_token:
            data["grant_type"] = "refresh_token"
            data["refresh_token"] = self.config.refresh_token
            
        data.update(self.config.extra_params)
        
        response = await client.post(self.config.token_url, data=data)
        response.raise_for_status()
        self._access_token = response.json()["access_token"]
        return self._access_token

class HttpClient:
    """
    An enhanced, instrumented async HTTP client.
    Handles auth, retries, normalization, and detailed logging.
    """
    def __init__(
        self, 
        name: str, 
        settings: ClientSettings
    ):
        self.name = name
        self.settings = settings
        
        # Normalization of base_url
        self.base_url = self._normalize_base_url(self.settings.base_url)

        self._clients_per_loop = weakref.WeakKeyDictionary()
        
        self._auth_manager = None
        if self.settings.auth:
            self._setup_auth(self.settings.auth)

    @property
    def client(self) -> httpx.AsyncClient:
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            try:
                loop = asyncio.get_event_loop()
            except RuntimeError:
                loop = None

        if loop is None:
            return self._build_httpx_client()

        if loop not in self._clients_per_loop:
            self._clients_per_loop[loop] = self._build_httpx_client()
            
        return self._clients_per_loop[loop]

    def _build_httpx_client(self) -> httpx.AsyncClient:
        client = httpx.AsyncClient(
            base_url=self.base_url, 
            timeout=self.settings.timeout,
            follow_redirects=self.settings.handle_redirects
        )
        if hasattr(self, "_basic_auth_tuple"):
            client.auth = httpx.BasicAuth(*self._basic_auth_tuple)
        return client

    def _normalize_base_url(self, url: str) -> str:
        if not url.endswith("/"):
            return url + "/"
        return url

    def _normalize_path(self, path: str) -> str:
        if path.startswith("/"):
            return path.lstrip("/")
        return path

    def _setup_auth(self, auth_config: AuthConfig):
        if auth_config.type == AuthType.BASIC:
            self._basic_auth_tuple = (auth_config.username, auth_config.password)
        elif auth_config.type == AuthType.API_KEY:
            # Handled in request()
            pass
        elif auth_config.type == AuthType.OAUTH2:
            self._auth_manager = OAuth2Manager(auth_config)

    async def authorize(
        self, 
        headers: Dict[str, str], 
        params: Dict[str, Any],
        cookies: Dict[str, str]
    ) -> None:
        if not self.settings.auth:
            return

        auth = self.settings.auth
        if auth.type == AuthType.API_KEY:
            if auth.in_ == ApiKeyLocation.HEADER:
                headers[auth.key] = auth.value
            elif auth.in_ == ApiKeyLocation.QUERY:
                params[auth.key] = auth.value
        elif auth.type == AuthType.OAUTH2 and self._auth_manager:
            token = await self._auth_manager.get_token(self.client)
            headers["Authorization"] = f"Bearer {token}"
        elif auth.type == AuthType.BASIC:
            pass  # Handled by client.auth
        else:
            raise ValueError(f"Unsupported auth type: {auth.type}")

    async def request(
        self, 
        method: str, 
        path: str, 
        *,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        cookies: Optional[Dict[str, str]] = None,
        json: Optional[Any] = None,
        data: Optional[Any] = None,
        content: Optional[Any] = None,
        files: Optional[Any] = None,
        timeout: Optional[Union[float, httpx.Timeout]] = None,
    ) -> httpx.Response:
        path = self._normalize_path(path)
        headers = headers or {}
        params = params or {}
        cookies = cookies or {}
        
        # Additional Auth Handling
        await self.authorize(headers, params, cookies)
        
        # Observability: Extract Request Payload
        req_content_type = headers.get("Content-Type")
        req_body_bytes = None
        try:
            if json is not None:
                import json as json_lib
                req_body_bytes = json_lib.dumps(json).encode("utf-8")
                if not req_content_type:
                    req_content_type = "application/json"
            elif content is not None and isinstance(content, (bytes, str)):
                req_body_bytes = content.encode("utf-8") if isinstance(content, str) else content
            elif data is not None and isinstance(data, (bytes, str)):
                req_body_bytes = data.encode("utf-8") if isinstance(data, str) else data
        except Exception as e:
            logger.warning(f"Failed to extract request payload for observability: {e}")

        start_time = datetime.datetime.now(UTC)

        # Retry Loop
        max_retries = self.settings.retry.max_retries
        current_try = 0
        
        while True:
            response = None
            try:
                response = await self.client.request(
                    method, 
                    path, 
                    params=params,
                    headers=headers,
                    json=json,
                    data=data,
                    content=content,
                    files=files,
                    cookies=cookies,
                    timeout=timeout
                )
                
                if response.is_error:
                    if await self._should_retry(response, current_try, max_retries):
                        current_try += 1
                        backoff = 2 ** current_try
                        logger.info(f"Retrying request to {self.name} ({current_try}/{max_retries}) after {backoff}s...")
                        await asyncio.sleep(backoff)
                        continue
                        
                    await self._log_error_response(response)
                    # Don't raise yet, handle Observability FAILED below
                
                # Observability: Handle Response
                resp_content_type = response.headers.get("Content-Type")
                resp_bytes = None
                try:
                    resp_bytes = await response.aread() 
                except Exception as e:
                    logger.warning(f"Failed to read response payload: {e}")

                end_time = datetime.datetime.now(UTC)
                state = "SUCCEEDED" if not response.is_error else "FAILED"
                
                await record_data_exchange(
                    DataExchangeEvent(
                        integration=self.settings.client_id,
                        channel="HTTP",
                        operation=f"{method} {path}",
                        remote_system=self.base_url,
                        address=f"{self.base_url}{path}",
                        occurred_at=start_time,
                        completed_at=end_time,
                        state=state,
                        attempt=current_try + 1,
                        http_method=method,
                        status_code=response.status_code,
                        response_payload=resp_bytes,
                        response_content_type=resp_content_type,
                        request_payload=req_body_bytes,
                        request_content_type=req_content_type
                    ),
                    correlation=None
                )

                response.raise_for_status()
                return response

            except httpx.HTTPStatusError:
                raise
            except (httpx.ConnectError, httpx.TimeoutException) as e:
                if current_try < max_retries:
                    current_try += 1
                    backoff = 2 ** current_try
                    logger.warning(f"Connection error to {self.name}, retrying ({current_try}/{max_retries}) context: {e}")
                    await asyncio.sleep(backoff)
                    continue
                
                # FINAL FAILURE - Observability: TIMEOUT/FAILED
                end_time = datetime.datetime.now(UTC)
                
                await record_data_exchange(
                    DataExchangeEvent(
                        integration=self.settings.client_id,
                        channel="HTTP",
                        operation=f"{method} {path}",
                        remote_system=self.base_url,
                        address=f"{self.base_url}{path}",
                        occurred_at=start_time,
                        completed_at=end_time,
                        state="TIMEOUT" if isinstance(e, httpx.TimeoutException) else "FAILED",
                        attempt=current_try + 1,
                        http_method=method,
                        attrs={"error": str(e)},
                        request_payload=req_body_bytes,
                        request_content_type=req_content_type
                    ),
                    correlation=None
                )
                
                raise e

    async def _should_retry(self, response: httpx.Response, current_try: int, max_retries: int) -> bool:
        if current_try >= max_retries:
            return False
            
        if response.status_code in [429, 502, 503, 504]:
            retry_possible = True
        else:
            retry_possible = False

        body = response.text
        for pattern in self.settings.retry.blacklist:
            if pattern in body:
                logger.debug(f"Retry blacklisted due to pattern '{pattern}' in response.")
                return False
        
        for pattern in self.settings.retry.whitelist:
            if pattern in body:
                logger.debug(f"Retry whitelisted due to pattern '{pattern}' in response.")
                return True
        
        return retry_possible

    async def _log_error_response(self, response: httpx.Response):
        try:
            body = response.text
            logger.error(
                f"HTTP Error {response.status_code} from {self.name}\n"
                f"URL: {response.url}\n"
                f"Response Body: {body[:1000]}{'...' if len(body) > 1000 else ''}"
            )
        except Exception as e:
            logger.error(f"Failed to log error response body: {e}")

    async def close(self):
        for client in self._clients_per_loop.values():
            await client.aclose()
        self._clients_per_loop.clear()
