"""AITracer main client."""

from __future__ import annotations

import atexit
import os
import threading
import time
import uuid
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Optional

from aitracer.pii import PIIDetector
from aitracer.queue import LogQueue
from aitracer.session import Session
from aitracer.trace import Trace
from aitracer.wrappers.openai_wrapper import wrap_openai_client
from aitracer.wrappers.anthropic_wrapper import wrap_anthropic_client
from aitracer.wrappers.gemini_wrapper import wrap_gemini_model

if TYPE_CHECKING:
    from openai import OpenAI
    from anthropic import Anthropic
    from google.generativeai import GenerativeModel


class AITracer:
    """
    AITracer client for monitoring AI/LLM applications.

    Usage:
        tracer = AITracer(api_key="at-xxxxxxxx", project="my-chatbot")
        client = tracer.wrap_openai(OpenAI())

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello!"}]
        )
    """

    _instance: Optional["AITracer"] = None
    _lock = threading.Lock()

    def __init__(
        self,
        api_key: Optional[str] = None,
        project: Optional[str] = None,
        *,
        user_id: Optional[str] = None,
        base_url: str = "https://api.aitracer.co",
        sync: bool = False,
        flush_on_exit: bool = True,
        batch_size: int = 10,
        flush_interval: float = 5.0,
        enabled: bool = True,
        # PII settings
        pii_detection: bool = False,
        pii_action: str = "mask",
        pii_types: Optional[list[str]] = None,
    ):
        """
        Initialize AITracer client.

        Args:
            api_key: AITracer API key. Can also be set via AITRACER_API_KEY env var.
            project: Project name or ID. Can also be set via AITRACER_PROJECT env var.
            user_id: Optional app user ID. Applied to all logs when not using sessions.
                     Can also be set via AITRACER_USER_ID env var.
            base_url: AITracer API base URL.
            sync: If True, send logs synchronously (useful for Lambda/serverless).
            flush_on_exit: If True, flush pending logs on program exit.
            batch_size: Number of logs to batch before sending.
            flush_interval: Seconds between automatic flushes.
            enabled: If False, disable all logging (useful for testing).
            pii_detection: If True, enable PII detection and masking.
            pii_action: Action for detected PII - "mask", "redact", "hash", "none".
            pii_types: List of PII types to detect. Default: ["email", "phone", "credit_card", "ssn"].
        """
        self.api_key = api_key or os.environ.get("AITRACER_API_KEY")
        self.project = project or os.environ.get("AITRACER_PROJECT")
        self.user_id = user_id or os.environ.get("AITRACER_USER_ID")
        self.base_url = base_url.rstrip("/")
        self.sync = sync
        self.enabled = enabled

        # PII detection settings
        self.pii_detection = pii_detection
        self.pii_action = pii_action
        self.pii_types = pii_types or ["email", "phone", "credit_card", "ssn"]
        self._pii_detector: Optional[PIIDetector] = None
        if self.pii_detection:
            self._pii_detector = PIIDetector(
                enabled_types=self.pii_types,
                action=self.pii_action,
            )

        if not self.api_key:
            raise ValueError(
                "API key is required. Set it via api_key parameter or AITRACER_API_KEY env var."
            )

        if not self.project:
            raise ValueError(
                "Project is required. Set it via project parameter or AITRACER_PROJECT env var."
            )

        # Log queue for async sending
        self._queue = LogQueue(
            api_key=self.api_key,
            base_url=self.base_url,
            batch_size=batch_size,
            flush_interval=flush_interval,
            sync=sync,
        )

        # Current trace context (thread-local)
        self._trace_context = threading.local()

        # Current session context (thread-local)
        self._session_context = threading.local()

        # Register flush on exit
        if flush_on_exit:
            atexit.register(self.flush)

        # Set as singleton instance
        with AITracer._lock:
            AITracer._instance = self

    @classmethod
    def get_instance(cls) -> Optional["AITracer"]:
        """Get the current AITracer instance."""
        return cls._instance

    def wrap_openai(self, client: "OpenAI") -> "OpenAI":
        """
        Wrap an OpenAI client to automatically log all API calls.

        Args:
            client: OpenAI client instance.

        Returns:
            Wrapped OpenAI client.
        """
        return wrap_openai_client(client, self)

    def wrap_anthropic(self, client: "Anthropic") -> "Anthropic":
        """
        Wrap an Anthropic client to automatically log all API calls.

        Args:
            client: Anthropic client instance.

        Returns:
            Wrapped Anthropic client.
        """
        return wrap_anthropic_client(client, self)

    def wrap_gemini(self, model: "GenerativeModel") -> "GenerativeModel":
        """
        Wrap a Google Gemini GenerativeModel to automatically log all API calls.

        Args:
            model: GenerativeModel instance.

        Returns:
            Wrapped GenerativeModel.

        Example:
            import google.generativeai as genai
            genai.configure(api_key="your-api-key")
            model = genai.GenerativeModel("gemini-1.5-flash")
            model = tracer.wrap_gemini(model)
        """
        return wrap_gemini_model(model, self)

    @contextmanager
    def trace(
        self,
        trace_id: Optional[str] = None,
        name: Optional[str] = None,
    ):
        """
        Create a trace context for grouping related API calls.

        Usage:
            with tracer.trace("user-query-123") as trace:
                response1 = client.chat.completions.create(...)
                response2 = client.chat.completions.create(...)
                trace.set_metadata({"user_id": "user-456"})

        Args:
            trace_id: Optional trace ID. Auto-generated if not provided.
            name: Optional name for the trace.

        Yields:
            Trace object for setting metadata.
        """
        trace = Trace(
            trace_id=trace_id or str(uuid.uuid4()),
            name=name,
        )

        # Store in thread-local context
        previous_trace = getattr(self._trace_context, "current", None)
        self._trace_context.current = trace

        try:
            yield trace
        finally:
            self._trace_context.current = previous_trace

    def get_current_trace(self) -> Optional[Trace]:
        """Get the current trace context."""
        return getattr(self._trace_context, "current", None)

    @contextmanager
    def session(
        self,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ):
        """
        Create a session context for tracking user interactions over time.

        A session groups multiple logs/traces that belong to the same user interaction flow.

        Usage:
            with tracer.session(user_id="user-123") as session:
                response = client.chat.completions.create(...)
                session.thumbs_up()  # Record positive feedback for last response

        Args:
            session_id: Optional session ID. Auto-generated if not provided.
            user_id: Optional user identifier.
            name: Optional name for the session.
            metadata: Optional metadata dictionary.

        Yields:
            Session object for recording events and feedback.
        """
        sess = Session(
            session_id=session_id or str(uuid.uuid4()),
            user_id=user_id,
            name=name,
            metadata=metadata or {},
            _tracer=self,
        )

        # Store in thread-local context
        previous_session = getattr(self._session_context, "current", None)
        self._session_context.current = sess

        # Start session on server
        if self.enabled:
            self._start_session(sess)

        try:
            yield sess
        finally:
            # End session on server
            if self.enabled:
                self._end_session(sess)
            self._session_context.current = previous_session

    def get_current_session(self) -> Optional[Session]:
        """Get the current session context."""
        return getattr(self._session_context, "current", None)

    def _start_session(self, session: Session) -> None:
        """Internal: Start a session on the server."""
        import httpx

        try:
            with httpx.Client() as client:
                client.post(
                    f"{self.base_url}/api/v1/sessions/",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "session_id": session.session_id,
                        "user_id": session.user_id,
                        "name": session.name,
                        "metadata": session.metadata,
                        "project": self.project,
                    },
                    timeout=10.0,
                )
        except Exception:
            pass  # Silently fail - don't break user code

    def _end_session(self, session: Session) -> None:
        """Internal: End a session on the server."""
        import httpx

        try:
            with httpx.Client() as client:
                client.patch(
                    f"{self.base_url}/api/v1/sessions/{session.session_id}/end/",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "metadata": session.metadata,
                    },
                    timeout=10.0,
                )
        except Exception:
            pass

    def _send_session_event(self, session_id: str, event: dict) -> None:
        """Internal: Send a session event to the server."""
        import httpx

        try:
            with httpx.Client() as client:
                client.post(
                    f"{self.base_url}/api/v1/sessions/{session_id}/events/",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json",
                    },
                    json=event,
                    timeout=10.0,
                )
        except Exception:
            pass

    def _send_feedback(self, session_id: str, feedback: dict) -> None:
        """Internal: Send user feedback to the server."""
        import httpx

        try:
            with httpx.Client() as client:
                client.post(
                    f"{self.base_url}/api/v1/sessions/{session_id}/feedback/",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json",
                    },
                    json=feedback,
                    timeout=10.0,
                )
        except Exception:
            pass

    def log(
        self,
        *,
        model: str,
        provider: str,
        input_data: Any,
        output_data: Any,
        input_tokens: int = 0,
        output_tokens: int = 0,
        latency_ms: int = 0,
        status: str = "success",
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
        trace_id: Optional[str] = None,
        span_id: Optional[str] = None,
        parent_span_id: Optional[str] = None,
    ) -> None:
        """
        Log an LLM API call.

        This is called automatically when using wrapped clients,
        but can also be called manually for custom logging.
        """
        if not self.enabled:
            return

        # Apply PII detection and masking if enabled
        processed_input = input_data
        processed_output = output_data
        if self._pii_detector:
            processed_input, _ = self._pii_detector.process_json(input_data)
            if output_data:
                processed_output, _ = self._pii_detector.process_json(output_data)

        # Get trace context
        current_trace = self.get_current_trace()
        if current_trace and not trace_id:
            trace_id = current_trace.trace_id
            if current_trace.metadata:
                metadata = {**(metadata or {}), **current_trace.metadata}

        # Auto-generate trace_id if not provided
        if not trace_id:
            trace_id = str(uuid.uuid4())

        # Get session context
        session_id = None
        app_user_id = self.user_id  # Use tracer-level user_id as default
        current_session = self.get_current_session()
        if current_session:
            session_id = current_session.session_id
            # Session user_id takes precedence over tracer-level user_id
            if current_session.user_id:
                app_user_id = current_session.user_id

        # Generate span_id
        generated_span_id = span_id or str(uuid.uuid4())

        log_entry = {
            "project": self.project,
            "model": model,
            "provider": provider,
            "input": processed_input,
            "output": processed_output,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "latency_ms": latency_ms,
            "status": status,
            "error_message": error_message,
            "metadata": metadata or {},
            "trace_id": trace_id,
            "span_id": generated_span_id,
            "parent_span_id": parent_span_id,
            "session_id": session_id,
            "app_user_id": app_user_id,
        }

        self._queue.add(log_entry)

        # Update session's last log ID for feedback association
        if current_session:
            current_session._set_last_log_id(generated_span_id)

    def flush(self) -> None:
        """Flush all pending logs immediately."""
        self._queue.flush()

    def shutdown(self) -> None:
        """Shutdown the client and flush pending logs."""
        self.flush()
        self._queue.shutdown()
