"""Persistence sinks for nanomon."""

from __future__ import annotations

import json
import logging
import os
import sqlite3
import time
import uuid
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from nanomon.errors import SinkWriteError

logger = logging.getLogger(__name__)


class Sink(ABC):
    """Abstract base class for run persistence sinks."""

    @abstractmethod
    def create_context(self, name: str | None = None, metadata: dict[str, Any] | None = None) -> str:
        """Create a run context and return its ID.

        Args:
            name: Optional name for the context
            metadata: Optional metadata for the context

        Returns:
            The ID of the created context

        Raises:
            SinkWriteError: If the operation fails
        """
        ...

    @abstractmethod
    def write_run(self, run_payload: dict[str, Any]) -> None:
        """Write a run payload to the sink (create).

        Args:
            run_payload: The run data to persist

        Raises:
            SinkWriteError: If the write operation fails
        """
        ...

    @abstractmethod
    def update_run(self, run_id: str, run_payload: dict[str, Any]) -> None:
        """Update an existing run in the sink.

        Args:
            run_id: The ID of the run to update
            run_payload: The run data to persist

        Raises:
            SinkWriteError: If the write operation fails
        """
        ...

    @abstractmethod
    def write_tool_call(self, run_id: str, tool_payload: dict[str, Any]) -> None:
        """Write a tool call record to the sink.

        Args:
            run_id: The ID of the run this tool call belongs to
            tool_payload: The tool call data to persist (name, method, startedAt, etc.)

        Raises:
            SinkWriteError: If the write operation fails
        """
        ...


class ApiSink(Sink):
    """Sink that writes runs to an HTTP API endpoint.

    Sends POST requests to {base_url}/runs with JSON payload.
    No authentication headers are included.
    """

    def __init__(
        self,
        base_url: str,
        *,
        timeout_s: float = 5.0,
        retries: int = 0,
    ) -> None:
        """Initialize the API sink.

        Args:
            base_url: Base URL for the API (e.g., "http://localhost:8000/api/v1")
            timeout_s: Request timeout in seconds
            retries: Number of retries for 5xx errors or network failures
        """
        self._base_url = base_url.rstrip("/")
        self._timeout_s = timeout_s
        self._retries = retries

        # Read auth headers from environment
        self._client_id = os.environ.get("NANOMON_CLIENT_ID")
        self._project_secret = os.environ.get("NANOMON_SECRET_ID")

    def _get_headers(self, method: str) -> dict[str, str]:
        """Get headers for a request, including auth for POST/PUT.

        Args:
            method: HTTP method (POST, PUT, GET, etc.)

        Returns:
            Dictionary of headers to include in the request
        """
        headers = {"Content-Type": "application/json"}

        if method in ("POST", "PUT") and self._client_id and self._project_secret:
            headers["X-CLIENT-ID"] = self._client_id
            headers["X-PROJECT-SECRET"] = self._project_secret

        return headers

    def _log_http_error(self, status_code: int, operation: str) -> None:
        """Log HTTP errors with appropriate messages based on status code."""
        if status_code == 401:
            logger.error(
                "Nanomon authentication failed. "
                "Please check your project credentials at https://app.nanomon.ai"
            )
        elif status_code == 403:
            logger.warning(
                "Nanomon quota exceeded for your subscription plan. "
                "Upgrade at https://app.nanomon.ai to continue tracking."
            )
        else:
            logger.error(f"Nanomon API error ({status_code}) during {operation}")

    def create_context(self, name: str | None = None, metadata: dict[str, Any] | None = None) -> str:
        """Create a run context via the API.

        Args:
            name: Optional name for the context
            metadata: Optional metadata for the context

        Returns:
            The ID of the created context (or a generated fallback ID on error)
        """
        url = f"{self._base_url}/contexts"

        try:
            payload = {"name": name, "metadata": metadata or {}}
            data = json.dumps(payload).encode("utf-8")
        except (TypeError, ValueError) as e:
            logger.warning(f"Nanomon failed to encode context payload: {e}")
            return str(uuid.uuid4())

        attempt = 0
        max_attempts = self._retries + 1

        while attempt < max_attempts:
            try:
                request = Request(
                    url,
                    data=data,
                    headers=self._get_headers("POST"),
                    method="POST",
                )

                with urlopen(request, timeout=self._timeout_s) as response:
                    status_code = response.getcode()
                    if 200 <= status_code < 300:
                        try:
                            response_data = json.loads(response.read().decode("utf-8"))
                            context_id = response_data["id"]
                            logger.debug(f"Successfully created context {context_id}")
                            return context_id
                        except (json.JSONDecodeError, KeyError) as e:
                            logger.warning(f"Nanomon failed to parse context response: {e}")
                            return str(uuid.uuid4())

            except HTTPError as e:
                status_code = e.code
                self._log_http_error(status_code, "create_context")
                # Don't retry on 4xx errors
                if 400 <= status_code < 500:
                    return str(uuid.uuid4())
                logger.debug(f"API returned {status_code}, attempt {attempt + 1}/{max_attempts}")

            except URLError as e:
                logger.warning(f"Nanomon network error: {e}")

            except TimeoutError:
                logger.warning("Nanomon request timeout")

            except Exception as e:
                logger.warning(f"Nanomon unexpected error: {e}")
                return str(uuid.uuid4())

            attempt += 1
            if attempt < max_attempts:
                sleep_time = 2 ** (attempt - 1)
                time.sleep(sleep_time)

        # All retries exhausted - return fallback ID
        logger.warning("Nanomon failed to create context after all retries, using local ID")
        return str(uuid.uuid4())

    def write_run(self, run_payload: dict[str, Any]) -> None:
        """Write a run to the API.

        Args:
            run_payload: The run data to send

        Note:
            This method never raises exceptions. Errors are logged and swallowed
            to ensure monitoring never crashes the user's flow.
        """
        url = f"{self._base_url}/runs"

        try:
            data = json.dumps(run_payload).encode("utf-8")
        except (TypeError, ValueError) as e:
            logger.warning(f"Nanomon failed to encode run payload: {e}")
            return

        attempt = 0
        max_attempts = self._retries + 1

        while attempt < max_attempts:
            try:
                request = Request(
                    url,
                    data=data,
                    headers=self._get_headers("POST"),
                    method="POST",
                )

                with urlopen(request, timeout=self._timeout_s) as response:
                    status_code = response.getcode()
                    if 200 <= status_code < 300:
                        logger.debug(f"Successfully wrote run {run_payload.get('id')} to API")
                        return

            except HTTPError as e:
                status_code = e.code
                self._log_http_error(status_code, "write_run")
                # Don't retry on 4xx errors
                if 400 <= status_code < 500:
                    return
                logger.debug(f"API returned {status_code}, attempt {attempt + 1}/{max_attempts}")

            except URLError as e:
                logger.warning(f"Nanomon network error: {e}")

            except TimeoutError:
                logger.warning("Nanomon request timeout")

            except Exception as e:
                logger.warning(f"Nanomon unexpected error: {e}")
                return

            attempt += 1
            if attempt < max_attempts:
                sleep_time = 2 ** (attempt - 1)
                time.sleep(sleep_time)

        # All retries exhausted - just log and return
        logger.warning("Nanomon failed to write run after all retries")

    def update_run(self, run_id: str, run_payload: dict[str, Any]) -> None:
        """Update an existing run in the API.

        Args:
            run_id: The ID of the run to update
            run_payload: The run data to send

        Note:
            This method never raises exceptions. Errors are logged and swallowed
            to ensure monitoring never crashes the user's flow.
        """
        url = f"{self._base_url}/runs/{run_id}"

        try:
            data = json.dumps(run_payload).encode("utf-8")
        except (TypeError, ValueError) as e:
            logger.warning(f"Nanomon failed to encode run payload: {e}")
            return

        attempt = 0
        max_attempts = self._retries + 1

        while attempt < max_attempts:
            try:
                request = Request(
                    url,
                    data=data,
                    headers=self._get_headers("PUT"),
                    method="PUT",
                )

                with urlopen(request, timeout=self._timeout_s) as response:
                    status_code = response.getcode()
                    if 200 <= status_code < 300:
                        logger.debug(f"Successfully updated run {run_id} in API")
                        return

            except HTTPError as e:
                status_code = e.code
                self._log_http_error(status_code, "update_run")
                # Don't retry on 4xx errors
                if 400 <= status_code < 500:
                    return
                logger.debug(f"API returned {status_code}, attempt {attempt + 1}/{max_attempts}")

            except URLError as e:
                logger.warning(f"Nanomon network error: {e}")

            except TimeoutError:
                logger.warning("Nanomon request timeout")

            except Exception as e:
                logger.warning(f"Nanomon unexpected error: {e}")
                return

            attempt += 1
            if attempt < max_attempts:
                sleep_time = 2 ** (attempt - 1)
                time.sleep(sleep_time)

        # All retries exhausted - just log and return
        logger.warning("Nanomon failed to update run after all retries")

    def write_tool_call(self, run_id: str, tool_payload: dict[str, Any]) -> None:
        """Write a tool call to the API.

        Args:
            run_id: The ID of the run this tool call belongs to
            tool_payload: The tool call data to send

        Note:
            This method never raises exceptions. Errors are logged and swallowed
            to ensure monitoring never crashes the user's flow.
        """
        url = f"{self._base_url}/runs/{run_id}/tools"

        try:
            data = json.dumps(tool_payload).encode("utf-8")
        except (TypeError, ValueError) as e:
            logger.warning(f"Nanomon failed to encode tool call payload: {e}")
            return

        attempt = 0
        max_attempts = self._retries + 1

        while attempt < max_attempts:
            try:
                request = Request(
                    url,
                    data=data,
                    headers=self._get_headers("POST"),
                    method="POST",
                )

                with urlopen(request, timeout=self._timeout_s) as response:
                    status_code = response.getcode()
                    if 200 <= status_code < 300:
                        logger.debug(f"Successfully wrote tool call {tool_payload.get('name')} for run {run_id}")
                        return

            except HTTPError as e:
                status_code = e.code
                self._log_http_error(status_code, "write_tool_call")
                # Don't retry on 4xx errors
                if 400 <= status_code < 500:
                    return
                logger.debug(f"API returned {status_code}, attempt {attempt + 1}/{max_attempts}")

            except URLError as e:
                logger.warning(f"Nanomon network error: {e}")

            except TimeoutError:
                logger.warning("Nanomon request timeout")

            except Exception as e:
                logger.warning(f"Nanomon unexpected error: {e}")
                return

            attempt += 1
            if attempt < max_attempts:
                sleep_time = 2 ** (attempt - 1)
                time.sleep(sleep_time)

        # All retries exhausted - just log and return
        logger.warning("Nanomon failed to write tool call after all retries")


class SqliteSink(Sink):
    """Sink that writes runs to a local SQLite database."""

    def __init__(self, db_path: str | Path) -> None:
        """Initialize the SQLite sink.

        Args:
            db_path: Path to the SQLite database file
        """
        self._db_path = Path(db_path)
        self._ensure_schema()

    def _ensure_schema(self) -> None:
        """Create the database schema if it doesn't exist."""
        with sqlite3.connect(self._db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS run_contexts (
                    id TEXT PRIMARY KEY,
                    name TEXT,
                    metadata TEXT DEFAULT '{}',
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS runs (
                    id TEXT PRIMARY KEY,
                    run_context_id TEXT,
                    started_at TEXT NOT NULL,
                    ended_at TEXT,
                    duration_ms INTEGER,
                    status TEXT NOT NULL,
                    error TEXT,
                    provider TEXT,
                    primary_model TEXT,
                    input_tokens INTEGER DEFAULT 0,
                    output_tokens INTEGER DEFAULT 0,
                    total_tokens INTEGER DEFAULT 0,
                    cost_usd REAL DEFAULT 0.0,
                    latency_ms INTEGER DEFAULT 0,
                    tags TEXT,
                    metadata TEXT,
                    by_model TEXT,
                    input TEXT,
                    output TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS tool_calls (
                    id TEXT PRIMARY KEY,
                    run_id TEXT NOT NULL,
                    name TEXT NOT NULL,
                    method TEXT,
                    started_at TEXT NOT NULL,
                    ended_at TEXT,
                    duration_ms INTEGER DEFAULT 0,
                    status TEXT NOT NULL,
                    error TEXT,
                    input TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (run_id) REFERENCES runs(id) ON DELETE CASCADE
                )
            """)
            conn.commit()

    def create_context(self, name: str | None = None, metadata: dict[str, Any] | None = None) -> str:
        """Create a run context in SQLite.

        Args:
            name: Optional name for the context
            metadata: Optional metadata for the context

        Returns:
            The ID of the created context

        Raises:
            SinkWriteError: If the database operation fails
        """
        try:
            context_id = str(uuid.uuid4())
            with sqlite3.connect(self._db_path) as conn:
                conn.execute(
                    "INSERT INTO run_contexts (id, name, metadata) VALUES (?, ?, ?)",
                    (context_id, name, json.dumps(metadata or {})),
                )
                conn.commit()
                logger.debug(f"Successfully created context {context_id} in SQLite")
                return context_id
        except sqlite3.Error as e:
            raise SinkWriteError(f"SQLite error: {e}") from e

    def write_run(self, run_payload: dict[str, Any]) -> None:
        """Write a run to the SQLite database.

        Args:
            run_payload: The run data to persist

        Raises:
            SinkWriteError: If the database operation fails
        """
        try:
            with sqlite3.connect(self._db_path) as conn:
                conn.execute(
                    """
                    INSERT OR REPLACE INTO runs (
                        id, run_context_id, started_at, ended_at, duration_ms, status, error,
                        provider, primary_model, input_tokens, output_tokens,
                        total_tokens, cost_usd, latency_ms, tags, metadata, by_model,
                        input, output
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        run_payload["id"],
                        run_payload.get("runContextId"),
                        run_payload["startedAt"],
                        run_payload.get("endedAt"),
                        run_payload.get("durationMs", 0),
                        run_payload["status"],
                        run_payload.get("error"),
                        run_payload.get("provider"),
                        run_payload.get("primaryModel"),
                        run_payload.get("inputTokens", 0),
                        run_payload.get("outputTokens", 0),
                        run_payload.get("totalTokens", 0),
                        run_payload.get("costUsd", 0.0),
                        run_payload.get("latencyMs", 0),
                        json.dumps(run_payload.get("tags", [])),
                        json.dumps(run_payload.get("metadata", {})),
                        json.dumps(run_payload.get("byModel", [])),
                        json.dumps(run_payload.get("input")) if run_payload.get("input") else None,
                        json.dumps(run_payload.get("output")) if run_payload.get("output") else None,
                    ),
                )
                conn.commit()
                logger.debug(f"Successfully wrote run {run_payload.get('id')} to SQLite")

        except sqlite3.Error as e:
            raise SinkWriteError(f"SQLite error: {e}") from e

    def update_run(self, run_id: str, run_payload: dict[str, Any]) -> None:
        """Update an existing run in SQLite (uses INSERT OR REPLACE).

        Args:
            run_id: The ID of the run to update
            run_payload: The run data to persist

        Raises:
            SinkWriteError: If the database operation fails
        """
        # SQLite uses INSERT OR REPLACE, so this is the same as write_run
        self.write_run(run_payload)

    def write_tool_call(self, run_id: str, tool_payload: dict[str, Any]) -> None:
        """Write a tool call to the SQLite database.

        Args:
            run_id: The ID of the run this tool call belongs to
            tool_payload: The tool call data to persist

        Raises:
            SinkWriteError: If the database operation fails
        """
        try:
            tool_id = tool_payload.get("id") or str(uuid.uuid4())
            started_at = tool_payload["startedAt"]
            ended_at = tool_payload.get("endedAt")

            # Calculate duration
            if ended_at:
                from datetime import datetime
                start_dt = datetime.fromisoformat(started_at.replace("Z", "+00:00"))
                end_dt = datetime.fromisoformat(ended_at.replace("Z", "+00:00"))
                duration_ms = int((end_dt - start_dt).total_seconds() * 1000)
            else:
                duration_ms = 0

            with sqlite3.connect(self._db_path) as conn:
                conn.execute(
                    """
                    INSERT INTO tool_calls (
                        id, run_id, name, method, started_at, ended_at,
                        duration_ms, status, error, input
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        tool_id,
                        run_id,
                        tool_payload["name"],
                        tool_payload.get("method"),
                        started_at,
                        ended_at,
                        duration_ms,
                        tool_payload["status"],
                        tool_payload.get("error"),
                        json.dumps(tool_payload.get("input")) if tool_payload.get("input") else None,
                    ),
                )
                conn.commit()
                logger.debug(f"Successfully wrote tool call {tool_payload['name']} for run {run_id}")

        except sqlite3.Error as e:
            raise SinkWriteError(f"SQLite error: {e}") from e


class MultiSink(Sink):
    """Composite sink that writes to multiple sinks.

    By default, raises an error if any sink fails.
    Can be configured to continue on failure.
    """

    def __init__(
        self,
        sinks: list[Sink],
        *,
        continue_on_failure: bool = False,
    ) -> None:
        """Initialize the multi-sink.

        Args:
            sinks: List of sinks to write to
            continue_on_failure: If True, continue writing to remaining sinks
                                 even if one fails. Raises the first error at the end.
        """
        self._sinks = sinks
        self._continue_on_failure = continue_on_failure

    def create_context(self, name: str | None = None, metadata: dict[str, Any] | None = None) -> str:
        """Create a context using the first sink and return its ID.

        Args:
            name: Optional name for the context
            metadata: Optional metadata for the context

        Returns:
            The ID of the created context (from the first sink)

        Raises:
            SinkWriteError: If the first sink fails
        """
        if not self._sinks:
            raise SinkWriteError("No sinks configured")
        # Use the first sink to create the context and get the ID
        return self._sinks[0].create_context(name=name, metadata=metadata)

    def write_run(self, run_payload: dict[str, Any]) -> None:
        """Write a run to all configured sinks.

        Args:
            run_payload: The run data to persist

        Raises:
            SinkWriteError: If any sink fails (behavior depends on continue_on_failure)
        """
        first_error: SinkWriteError | None = None

        for sink in self._sinks:
            try:
                sink.write_run(run_payload)
            except SinkWriteError as e:
                if first_error is None:
                    first_error = e

                if not self._continue_on_failure:
                    raise

                logger.warning(f"Sink {sink.__class__.__name__} failed: {e}")

        if first_error is not None:
            raise first_error

    def update_run(self, run_id: str, run_payload: dict[str, Any]) -> None:
        """Update a run in all configured sinks.

        Args:
            run_id: The ID of the run to update
            run_payload: The run data to persist

        Raises:
            SinkWriteError: If any sink fails (behavior depends on continue_on_failure)
        """
        first_error: SinkWriteError | None = None

        for sink in self._sinks:
            try:
                sink.update_run(run_id, run_payload)
            except SinkWriteError as e:
                if first_error is None:
                    first_error = e

                if not self._continue_on_failure:
                    raise

                logger.warning(f"Sink {sink.__class__.__name__} failed: {e}")

        if first_error is not None:
            raise first_error

    def write_tool_call(self, run_id: str, tool_payload: dict[str, Any]) -> None:
        """Write a tool call to all configured sinks.

        Args:
            run_id: The ID of the run this tool call belongs to
            tool_payload: The tool call data to persist

        Raises:
            SinkWriteError: If any sink fails (behavior depends on continue_on_failure)
        """
        first_error: SinkWriteError | None = None

        for sink in self._sinks:
            try:
                sink.write_tool_call(run_id, tool_payload)
            except SinkWriteError as e:
                if first_error is None:
                    first_error = e

                if not self._continue_on_failure:
                    raise

                logger.warning(f"Sink {sink.__class__.__name__} failed: {e}")

        if first_error is not None:
            raise first_error
