"""
Job lookup and management CLI commands.

Provides commands for querying job status and history by dataset or job ID.
"""

from __future__ import annotations

import os
import boto3
import typer
from rich.console import Console
from rich.table import Table

app = typer.Typer(help="Job lookup and management commands")


class JobsCommands:
    """
    Job lookup and management commands.

    Provides query and display functionality for jobs across all stages.
    """

    def __init__(self):
        """Initialize jobs commands with AWS clients."""
        # Create session with profile if available
        session = boto3.Session()
        if os.getenv("AWS_PROFILE"):
            session = boto3.Session(profile_name=os.getenv("AWS_PROFILE"))

        self.dynamodb = session.resource("dynamodb")
        self.table_name = os.getenv("DYNAMODB_TABLE_NAME", "czbg-jobs-dev")
        self.table = self.dynamodb.Table(self.table_name)
        self.console = Console()

    def list(
        self,
        dataset: str | None = None,
        status: str | None = None,
        stage: str | None = None,
        limit: int = 20,
    ) -> None:
        """
        List jobs with optional filtering.

        Args:
            dataset: Filter by dataset name
            status: Filter by status (queued, running, completed, failed)
            stage: Filter by job type/stage
            limit: Maximum number of jobs to return
        """
        # Query jobs
        if dataset:
            jobs = self._query_by_dataset(dataset)
        else:
            jobs = self._scan_all_jobs(limit)

        if not jobs:
            self.console.print("[yellow]No jobs found[/yellow]")
            return

        # Apply filters
        if status:
            jobs = [j for j in jobs if j.get("status") == status]
        if stage:
            jobs = [j for j in jobs if j.get("job_type") == stage]

        # Limit results
        jobs = jobs[:limit]

        if not jobs:
            self.console.print("[yellow]No jobs match filters[/yellow]")
            return

        # Display table
        self._display_jobs_table(jobs)

    def describe(self, dataset_or_job_id: str, verbose: bool = False) -> None:
        """
        Describe job(s) by dataset name or job ID.

        Args:
            dataset_or_job_id: Dataset name or job ID to lookup
            verbose: Show detailed information including configs
        """
        # Try dataset_name query first
        jobs = self._query_by_dataset(dataset_or_job_id)

        if not jobs:
            # Fallback to job_id lookup
            job = self._get_by_job_id(dataset_or_job_id)
            if job:
                jobs = [job]

        if not jobs:
            self.console.print(f"[red]No jobs found for: {dataset_or_job_id}[/red]")
            return

        # Display jobs grouped by dataset
        dataset_name = jobs[0].get("dataset_name", "unknown")
        self._display_pipeline_summary(dataset_name, jobs, verbose)

    def _query_by_dataset(self, dataset_name: str) -> list[dict]:
        """
        Query jobs by dataset name using GSI.

        Args:
            dataset_name: Dataset identifier

        Returns:
            List of job dictionaries
        """
        try:
            response = self.table.query(
                IndexName="dataset_name-created_at-index",
                KeyConditionExpression="dataset_name = :dn",
                ExpressionAttributeValues={":dn": dataset_name},
            )
            return response.get("Items", [])
        except Exception as e:
            self.console.print(f"[yellow]Warning: Failed to query by dataset: {e}[/yellow]")
            return []

    def _get_by_job_id(self, job_id: str) -> dict | None:
        """
        Get job by job_id.

        Args:
            job_id: Job identifier

        Returns:
            Job dictionary or None
        """
        try:
            response = self.table.get_item(Key={"job_id": job_id})
            return response.get("Item")
        except Exception as e:
            self.console.print(f"[yellow]Warning: Failed to get job: {e}[/yellow]")
            return None

    def _scan_all_jobs(self, limit: int) -> list[dict]:
        """
        Scan all jobs from table.

        Args:
            limit: Maximum number of jobs to return

        Returns:
            List of job dictionaries
        """
        try:
            response = self.table.scan(Limit=limit * 2)  # Get extra for filtering
            jobs = response.get("Items", [])

            # Sort by created_at descending (most recent first)
            jobs_sorted = sorted(
                jobs, key=lambda j: j.get("created_at", ""), reverse=True
            )

            return jobs_sorted
        except Exception as e:
            self.console.print(f"[yellow]Warning: Failed to scan jobs: {e}[/yellow]")
            return []

    def _display_jobs_table(self, jobs: list[dict]) -> None:
        """
        Display jobs in a table format.

        Args:
            jobs: List of job dictionaries
        """
        # Display header
        self.console.print()
        self.console.print(f"[bold]Found {len(jobs)} job(s)[/bold]")
        self.console.print()

        # Build table
        table = Table(show_header=True, header_style="bold")
        table.add_column("Dataset", style="cyan")
        table.add_column("Stage", style="magenta")
        table.add_column("Status")
        table.add_column("Created", style="dim")
        table.add_column("Duration")

        for job in jobs:
            dataset_name = job.get("dataset_name", "unknown")
            job_type = job.get("job_type", "unknown")
            status = job.get("status", "unknown")
            created_at = job.get("created_at", "")
            completed_at = job.get("completed_at", "")

            # Format created_at as date + time
            created_formatted = created_at.replace("T", " ")[:19] if created_at else "-"

            # Calculate duration
            duration = "-"
            if completed_at and created_at:
                try:
                    from datetime import datetime

                    start = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
                    end = datetime.fromisoformat(completed_at.replace("Z", "+00:00"))
                    duration_seconds = int((end - start).total_seconds())
                    duration = f"{duration_seconds}s"
                except:
                    pass

            # Format status with color
            status_formatted = self._format_status(status)

            table.add_row(
                dataset_name, job_type, status_formatted, created_formatted, duration
            )

        self.console.print(table)
        self.console.print()

    def _display_pipeline_summary(
        self, dataset_name: str, jobs: list[dict], verbose: bool
    ) -> None:
        """
        Display pipeline summary for dataset.

        Args:
            dataset_name: Dataset identifier
            jobs: List of job dictionaries
            verbose: Show detailed information
        """
        # Sort jobs by created_at
        jobs_sorted = sorted(jobs, key=lambda j: j.get("created_at", ""))

        # Display header
        self.console.print()
        self.console.print(f"[bold]Dataset:[/bold] {dataset_name}")
        if jobs_sorted:
            created_at = jobs_sorted[0].get("created_at", "unknown")
            self.console.print(f"[dim]Created: {created_at}[/dim]")
        self.console.print()

        # Build pipeline table
        table = Table(show_header=True, header_style="bold")
        table.add_column("Stage", style="cyan")
        table.add_column("Job ID", style="dim")
        table.add_column("Status")
        table.add_column("Duration")
        table.add_column("Completed")

        for job in jobs_sorted:
            job_type = job.get("job_type", "unknown")
            job_id = job.get("job_id", "unknown")
            status = job.get("status", "unknown")
            created_at = job.get("created_at", "")
            completed_at = job.get("completed_at", "")

            # Calculate duration
            duration = "-"
            if completed_at and created_at:
                try:
                    from datetime import datetime

                    start = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
                    end = datetime.fromisoformat(completed_at.replace("Z", "+00:00"))
                    duration_seconds = int((end - start).total_seconds())
                    duration = f"{duration_seconds}s"
                except:
                    pass

            # Format status with color
            status_formatted = self._format_status(status)

            # Format completed time
            completed_formatted = completed_at.split("T")[1][:8] if completed_at else "-"

            # Truncate job_id for display
            job_id_short = job_id[:8] + "..." if len(job_id) > 8 else job_id

            table.add_row(
                job_type, job_id_short, status_formatted, duration, completed_formatted
            )

        self.console.print(table)
        self.console.print()

        # Show processing breadcrumbs and errors for each job
        for job in jobs_sorted:
            # Show breadcrumbs if available
            if job.get("processing_stages"):
                self._display_processing_stages(job)

            # Show error details if failed
            if job.get("status") == "failed" and job.get("error_details"):
                self._display_error_details(job)

            # Show results for completed jobs
            if job.get("status") == "completed" and job.get("result"):
                self._display_job_result(job)

        # Show config paths if verbose
        if verbose:
            self._display_config_paths(jobs_sorted)

    def _format_status(self, status: str) -> str:
        """
        Format status with color.

        Args:
            status: Status string

        Returns:
            Formatted status string with Rich markup
        """
        status_colors = {
            "queued": "[yellow]⏳ Queued[/yellow]",
            "running": "[blue]▶ Running[/blue]",
            "completed": "[green]✓ Completed[/green]",
            "failed": "[red]✗ Failed[/red]",
        }
        return status_colors.get(status, f"[dim]{status}[/dim]")

    def _display_job_result(self, job: dict) -> None:
        """
        Display job result details.

        Args:
            job: Job dictionary with result
        """
        job_type = job.get("job_type", "unknown")
        result = job.get("result", {})

        self.console.print(f"[bold]{job_type.capitalize()} Stage Results:[/bold]")

        # Display key result fields
        for key, value in result.items():
            if key not in ("per_month_statistics", "peak_month"):
                self.console.print(f"  {key}: {value}")

        # Display per-month statistics if available
        if "per_month_statistics" in result and result["per_month_statistics"]:
            self._display_per_month_stats(result)

        self.console.print()

    def _display_per_month_stats(self, result: dict) -> None:
        """
        Display per-month statistics.

        Args:
            result: Job result with per_month_statistics
        """
        per_month_stats = result.get("per_month_statistics", {})
        peak_month = result.get("peak_month", {})

        self.console.print()
        self.console.print("[bold]Per-Month Line Item Counts:[/bold]")

        sorted_months = sorted(per_month_stats.keys())

        for month_key in sorted_months:
            stats = per_month_stats[month_key]
            provider = stats["provider"]
            month = stats["month"]
            count = stats["record_count"]
            size_mb = stats["file_size_mb"]

            is_peak = peak_month.get("key") == month_key

            if is_peak:
                self.console.print(
                    f"  [bold yellow]⭐ {provider}-{month}:[/bold yellow] "
                    f"[bold yellow]{count:,} records[/bold yellow] "
                    f"[dim]({size_mb} MB)[/dim] "
                    f"[bold yellow]← PEAK MONTH[/bold yellow]"
                )
            else:
                self.console.print(
                    f"  [cyan]{provider}-{month}:[/cyan] "
                    f"{count:,} records [dim]({size_mb} MB)[/dim]"
                )

    def _display_processing_stages(self, job: dict) -> None:
        """
        Display processing breadcrumbs.

        Args:
            job: Job dictionary with processing_stages
        """
        job_type = job.get("job_type", "unknown")
        processing_stages = job.get("processing_stages", [])

        if not processing_stages:
            return

        self.console.print(f"[bold]{job_type.capitalize()} Processing History:[/bold]")

        for stage in processing_stages:
            stage_name = stage.get("stage", "unknown")
            stage_status = stage.get("status", "unknown")
            timestamp = stage.get("timestamp", "")

            # Format timestamp (show time only)
            time_formatted = timestamp.split("T")[1][:8] if "T" in timestamp else timestamp

            # Icon based on status
            icon = "✓" if stage_status == "completed" else "▶" if stage_status == "in_progress" else "⏳"

            self.console.print(f"  {icon} {stage_name} [dim]({time_formatted})[/dim]")

        self.console.print()

    def _display_error_details(self, job: dict) -> None:
        """
        Display error details with suggestions.

        Args:
            job: Job dictionary with error_details
        """
        job_type = job.get("job_type", "unknown")
        error_details = job.get("error_details", {})

        if not error_details:
            return

        self.console.print(f"[bold red]{job_type.capitalize()} Error Details:[/bold red]")

        # Display error information
        stage = error_details.get("stage", "unknown")
        error_type = error_details.get("error_type", "Error")
        message = error_details.get("message", "No message")

        self.console.print(f"  Stage: {stage}")
        self.console.print(f"  Type: {error_type}")
        self.console.print(f"  Message: {message}")

        # Display field if available
        if "field" in error_details and error_details["field"]:
            self.console.print(f"  Field: {error_details['field']}")

        # Display suggestion if available
        if "suggestion" in error_details and error_details["suggestion"]:
            self.console.print()
            self.console.print(f"[yellow]💡 Suggestion:[/yellow]")
            self.console.print(f"  {error_details['suggestion']}")

        self.console.print()

    def _display_config_paths(self, jobs: list[dict]) -> None:
        """
        Display config S3 paths.

        Args:
            jobs: List of job dictionaries
        """
        self.console.print("[bold]Config Files:[/bold]")
        for job in jobs:
            job_type = job.get("job_type", "unknown")
            config_s3_path = job.get("config_s3_path", "unknown")
            self.console.print(f"  {job_type}: {config_s3_path}")
        self.console.print()


# CLI Commands
@app.command()
def list(
    dataset: str = typer.Option(None, "--dataset", "-d", help="Filter by dataset name"),
    status: str = typer.Option(None, "--status", "-s", help="Filter by status (queued, running, completed, failed)"),
    stage: str = typer.Option(None, "--stage", help="Filter by stage/job type (generate, k8s, metadata, usage, discount, upload)"),
    limit: int = typer.Option(20, "--limit", "-n", help="Maximum number of jobs to return"),
):
    """
    List jobs with optional filtering.

    Examples:
        czbg jobs list
        czbg jobs list --dataset 2025-12-19_15-14-46
        czbg jobs list --status failed
        czbg jobs list --stage k8s --limit 10
    """
    commands = JobsCommands()
    commands.list(dataset=dataset, status=status, stage=stage, limit=limit)


@app.command()
def describe(
    id: str = typer.Argument(..., help="Dataset name or job ID to lookup"),
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Show detailed information including configs"),
):
    """
    Describe job(s) by dataset name or job ID.

    Examples:
        czbg jobs describe 2025-12-19_15-14-46
        czbg jobs describe 83a0a7d0-c155-480d-a92a-f6f8cdf34f2f --verbose
    """
    commands = JobsCommands()
    commands.describe(id, verbose)
