"""Cost tracking integration with queue and metrics.

Provides integration between cost tracking and the job queue system
to automatically track costs for jobs.

Example:
    >>> from oclawma.costs.integration import CostsQueueIntegration
    >>> integration = CostsQueueIntegration()
    >>> integration.install(queue)
"""

from __future__ import annotations

import time
from contextlib import suppress
from typing import Any, Callable

from oclawma.costs import CostStore, CostTracker
from oclawma.metrics import MetricsCollector
from oclawma.queue import Job, JobQueue


class CostsQueueIntegration:
    """Integrates cost tracking with the job queue.

    Automatically tracks costs for jobs including:
    - API costs from job execution
    - Compute time
    - Storage usage

    Attributes:
        tracker: CostTracker instance
        store: CostStore instance
        original_callback: Original job callback (if any)
    """

    def __init__(self, store: CostStore | None = None) -> None:
        """Initialize integration.

        Args:
            store: CostStore instance (creates default if None)
        """
        self.store = store or CostStore()
        self.tracker = CostTracker(self.store)
        self._original_callback: Callable[[Job], None] | None = None
        self._job_start_times: dict[int, float] = {}

    def install(self, queue: JobQueue) -> None:
        """Install cost tracking hooks into a queue.

        Args:
            queue: JobQueue to integrate with
        """
        # Store original callback if exists
        if hasattr(queue, "_job_callback"):
            self._original_callback = queue._job_callback

        # Set up our callback
        queue._job_callback = self._on_job_complete

        # Hook into dequeue to track start times
        original_dequeue = queue.dequeue

        def tracked_dequeue(*args: Any, **kwargs: Any) -> Job | None:
            job = original_dequeue(*args, **kwargs)
            if job and job.id is not None:
                self._on_job_start(job)
            return job

        queue.dequeue = tracked_dequeue  # type: ignore

    def _on_job_start(self, job: Job) -> None:
        """Handle job start.

        Args:
            job: Job that started
        """
        if job.id is None:
            return

        job_id = f"job-{job.id}"
        self._job_start_times[job.id] = time.perf_counter()

        # Start tracking
        with suppress(Exception):
            self.tracker.start_job(
                job_id=job_id,
                project_id=job.payload.get("project_id"),
                tenant_id=job.tenant_id,
                metadata={
                    "queue_job_id": job.id,
                    "job_type": job.job_type,
                    "priority": (
                        job.priority.value if hasattr(job.priority, "value") else job.priority
                    ),
                },
            )

    def _on_job_complete(self, job: Job) -> None:
        """Handle job completion.

        Args:
            job: Job that completed
        """
        if job.id is None:
            if self._original_callback:
                self._original_callback(job)
            return

        job_id = f"job-{job.id}"

        # Calculate compute time
        start_time = self._job_start_times.pop(job.id, None)
        if start_time:
            compute_time = time.perf_counter() - start_time
        else:
            # Fallback: use job timestamps
            if job.completed_at and job.started_at:
                compute_time = (job.completed_at - job.started_at).total_seconds()
            else:
                compute_time = 0.0

        # Update job cost with compute time
        try:
            active_job = self.tracker._active_jobs.get(job_id)
            if active_job:
                active_job.compute_time_seconds = compute_time
                self.tracker.end_job(job_id)
        except Exception:
            # Don't let cost tracking errors break the queue
            pass

        # Call original callback
        if self._original_callback:
            with suppress(Exception):
                self._original_callback(job)

    def record_api_cost(
        self,
        job_id: int,
        provider: str,
        model: str,
        cost: float,
        tokens_input: int = 0,
        tokens_output: int = 0,
    ) -> None:
        """Record an API cost for a running job.

        This can be called from within job execution to track API costs.

        Args:
            job_id: Queue job ID
            provider: API provider name
            model: Model name
            cost: Cost in USD
            tokens_input: Input token count
            tokens_output: Output token count
        """
        cost_job_id = f"job-{job_id}"
        with suppress(Exception):
            self.tracker.record_api_cost(
                job_id=cost_job_id,
                provider=provider,
                model=model,
                cost=cost,
                tokens_input=tokens_input,
                tokens_output=tokens_output,
            )

    def close(self) -> None:
        """Close the integration and release resources."""
        self.tracker.close()


class CostsMetricsExporter:
    """Exports cost data to metrics collector.

    Provides cost-related metrics for Prometheus scraping.

    Attributes:
        store: CostStore instance
        collector: MetricsCollector instance
    """

    def __init__(
        self,
        store: CostStore | None = None,
        collector: MetricsCollector | None = None,
    ) -> None:
        """Initialize exporter.

        Args:
            store: CostStore instance
            collector: MetricsCollector instance
        """
        self.store = store or CostStore()
        self.collector = collector

    def export_daily_stats(self) -> dict[str, Any]:
        """Export daily cost statistics.

        Returns:
            Dictionary with daily statistics
        """
        from datetime import datetime

        from oclawma.costs import CostTracker, ReportPeriod

        tracker = CostTracker(self.store)

        try:
            report = tracker.generate_report(ReportPeriod.DAILY)

            stats = {
                "total_cost": report.total_cost,
                "total_jobs": report.total_jobs,
                "costs_by_project": report.costs_by_project,
                "costs_by_provider": report.costs_by_provider,
                "timestamp": datetime.utcnow().isoformat(),
            }

            # Update metrics collector if available
            if self.collector:
                self._update_metrics(report)

            return stats
        finally:
            tracker.close()

    def _update_metrics(self, report: Any) -> None:
        """Update metrics collector with cost data.

        Args:
            report: CostReport instance
        """
        # This would update custom metrics if the collector supports them
        # For now, we just store the values in the snapshot
        pass

    def close(self) -> None:
        """Close the exporter."""
        self.store.close()
