import re
import shutil
from pathlib import Path
from typing import List, Optional

import polars as pl
import typer
from rich.console import Console
from rich.table import Table

# Legacy reference/usage-reference imports removed - replaced by unified data commands
from src.cli.analysis import app as analysis_app
from src.cli.core.upload import app as upload_app
from src.cli.data import app as data_app
# Temporarily disabled due to legacy pipeline dependencies
# from cli.pipeline import app as pipeline_app
from src.cli.profiles import app as profiles_app
# from cli.telemetry import app as telemetry_app  # Temporarily disabled
from src.exploration import exploration
from src.generate.cli.generate_cli import app as generate_app
from src.k8s.cli.k8s_cli import k8s_app
from src.metadata.cli.metadata_cli import metadata_app
from src.usage.cli.usage_cli import usage_app
from src.discount.cli.discount_cli import discount_app
from src.services.s3_data_service import S3DataService

app = typer.Typer(help="CloudZero Billing Generator CLI")
# Temporarily disabled due to legacy pipeline dependencies
# app.add_typer(pipeline_app, name="pipeline")
# app.add_typer(telemetry_app, name="telemetry")  # Temporarily disabled
app.add_typer(profiles_app, name="profiles")
app.add_typer(analysis_app, name="analyze")
app.add_typer(data_app, name="data")
app.add_typer(upload_app, name="upload")
app.add_typer(generate_app, name="generate")
app.add_typer(k8s_app, name="k8s")
app.add_typer(metadata_app, name="metadata")
app.add_typer(usage_app, name="usage")
app.add_typer(discount_app, name="discount")
# Legacy reference/usage-reference commands removed - replaced by unified data commands

# Note: upload_app now uses S3-based approach (no workspace validation needed)

console = Console()

DEFAULT_CONFIG_YAML = """# config.yaml

org_name: {org}
api_token: ${{CLOUDBILLGEN_{env_suffix}_API_TOKEN}}
# api_url: https://api.demo.cloudzero.com/v1/ingest
connections:
  AWS:
    - your-aws-connection-id-here
  GCP:
    - your-gcp-connection-id-here
  Azure:
    - your-azure-connection-id-here

# Upload Configuration
upload:
  # Generated data upload settings
  generated:
    # Explicit data source selection
    data_source:
      run: "latest" # "latest" or specific run ID like "20241201_143022"
      stage: "generate" # Stage to upload (generate, validate, etc.)

    # Upload behavior settings (enforced maximums)
    settings:
      concurrency: 6 # Max: 20, Min: 1
      max_retries: 5 # Max: 10, Min: 1
      timeout: 60.0 # Max: 300.0, Min: 10.0
      dry_run: true
      resume: false

  # Ad-hoc CSV upload settings (FIXED VALUES - user changes ignored)
  adhoc:
    # These values are ENFORCED regardless of user changes
    input_dir: "adhoc/input" # FIXED: always adhoc/input/
    mapping_file: "adhoc/mapping.csv" # FIXED: always adhoc/mapping.csv
    output_dir: "adhoc/output" # FIXED: always adhoc/output/
    concurrency: 4 # FIXED: always 4
    max_retries: 3 # FIXED: always 3
    timeout: 30.0 # FIXED: always 30.0
    dry_run: false # FIXED: always false
"""

DEFAULT_PIPELINE_YAML = """# pipeline.jobs.yaml

stages:
  generate:
    - reference_data: data/reference.csv
      start_date: 2025-04-01
      end_date: 2025-04-30
      total_cost: 1234.56
      resource_count: 100
      filters:
        - cloud_provider: [AWS]
          variance_profile: "business_hours_daily"  # Realistic business hours variance patterns
          anomaly_profile: "quarterly_cost_optimization"  # Quarterly cost optimization cycles
          service: [EC2]
      

    - reference_data: data/reference.csv
      start_date: 2025-04-01
      end_date: 2025-04-30
      total_cost: 5678.90
      resource_count: 200
      filters:
        - cloud_provider: [GCP]
          variance_profile: "analytics_workload"  # Analytics burst processing patterns
          # No anomaly profile for this job - demonstrates mixed configuration
          service: [BigQuery]

  # Simplified metadata allocation
  metadata:
    majority_region: "us-east-1"        # 70% allocation to us-east-1
    majority_environment: "production"  # 70% allocation to production
    base_allocation: "products"         # Use products profile for allocation

  # Usage pricing calculation (auto-executes when --stage usage is run)
  usage:
    # No execution control flags - command-driven only
"""


DEFAULT_ENV = """# .env - Add your CloudZero API token here
# Get your API token from: https://app.cloudzero.com/organization/api-keys

CLOUDBILLGEN_{env_suffix}_API_TOKEN=your-cloudzero-api-token-here
"""


@app.command()
def init(
    force: bool = typer.Option(False, help="Overwrite existing files if they exist"),
):
    """Initialize ~/.czbg workspace with default structure and templates."""
    workspace_dir = Path.home() / ".czbg"
    
    # Create main directories
    workspace_dir.mkdir(exist_ok=True)
    (workspace_dir / "logs").mkdir(exist_ok=True)
    (workspace_dir / "configs").mkdir(exist_ok=True)
    (workspace_dir / "configs" / "generate").mkdir(exist_ok=True)
    (workspace_dir / "configs" / "upload").mkdir(exist_ok=True)
    (workspace_dir / "configs" / "k8s").mkdir(exist_ok=True)
    (workspace_dir / "configs" / "metadata").mkdir(exist_ok=True)
    (workspace_dir / "configs" / "usage").mkdir(exist_ok=True)
    (workspace_dir / "configs" / "discount").mkdir(exist_ok=True)
    
    # Template paths
    env_template_path = workspace_dir / ".env.template"
    generate_template_path = workspace_dir / "configs" / "generate" / "example-config.yaml"
    upload_template_path = workspace_dir / "configs" / "upload" / "upload.template.yaml"
    k8s_template_path = workspace_dir / "configs" / "k8s" / "k8s.template.yaml"
    metadata_template_path = workspace_dir / "configs" / "metadata" / "metadata.template.yaml"
    usage_template_path = workspace_dir / "configs" / "usage" / "usage.template.yaml"
    discount_template_path = workspace_dir / "configs" / "discount" / "discount.template.yaml"
    
    # Check if templates already exist (only warn, don't exit unless force is needed)
    templates_exist = env_template_path.exists()
    if templates_exist and not force:
        typer.echo(f"✔ Workspace already exists: {workspace_dir}")
        typer.echo("  Use --force to recreate template files")
        return
    
    # Create .env.template
    env_template_content = """# CloudZero API Tokens for All Organizations
# Copy this file to .env and add your actual tokens

# Example org tokens (replace with your actual org names and tokens):
# CLOUDZERO_ACME_CORP_API_TOKEN=your-acme-token-here
# CLOUDZERO_BETA_COMPANY_API_TOKEN=your-beta-token-here
# CLOUDZERO_DEMO_ORG_API_TOKEN=your-demo-token-here

# Get your tokens from: https://app.cloudzero.com/organization/api-keys

# S3 Configuration for Reference Data and Dataset Storage
# AWS credentials are handled automatically via boto3 default credential chain
# Supports: AWS SSO (aws sso login), environment variables, IAM roles, etc.
# No need to configure AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY in .env
AWS_DEFAULT_REGION=us-east-1
# AWS_PROFILE=your-aws-profile-name  # Optional: specify AWS profile for SSO

# S3 Bucket and Paths
CZBILLGEN_S3_BUCKET=your-czbillgen-bucket-name
CZBILLGEN_S3_REFERENCE_PATH=czbillgen/reference_data

# Generation Performance Settings
CZBILLGEN_ACCUMULATION_MEMORY_MB=500  # Memory limit for data accumulation before upload (MB)
CZBILLGEN_MEMORY_LOG_LEVEL=INFO       # Memory logging level: DEBUG, INFO, WARN, ERROR
"""
    
    # Create generate template
    generate_template_content = """# CloudZero Billing Generator - Generation Configuration Template
# Copy this file and customize for your generation needs

# TODO: Review and calibrate total_cost using 'czbg analyze config'
# TODO: Customize services, usage_families, and resource_types for your demo scenario

generation_jobs:
  - name: "Demo Multi-Cloud Generation"
    start_date: "2024-09-01"
    end_date: "2025-12-31"
    total_cost: 10000000.0  # $10M total across all providers
    resource_count: 25
    filters:
      - cloud_provider: AWS
        variance_profile: linear_growth
        anomaly_profile: exponential_burst
        services:
          - name: AmazonEC2
            usage_families:
              - name: Compute Instance
                resource_types:
                  - name: instance
                    operations:
                      - name: RunInstances
                        usage_types:
                          - BoxUsage:c5.4xlarge
                          - BoxUsage:r5.4xlarge
                          - BoxUsage:t3.medium
              - name: Storage
                resource_types:
                  - name: volume
                    operations:
                      - name: CreateVolume-Gp3
                        usage_types:
                          - EBS:VolumeUsage.gp3
          - name: AmazonS3
            usage_families:
              - name: Storage
                resource_types:
                  - name: bucket
                    operations:
                      - name: StandardStorage
                        usage_types:
                          - StandardStorage
                          - StandardIAStorage

      - cloud_provider: Azure
        variance_profile: steady_state
        anomaly_profile: ""
        services:
          - name: microsoft.compute
            usage_families:
              - name: Compute
                resource_types:
                  - name: virtualmachines
                    operations:
                      - name: Compute
                        usage_types:
                          - Standard_D4s_v3
                          - Standard_B2ms

      - cloud_provider: GCP
        variance_profile: periodic_surge
        anomaly_profile: ""
        services:
          - name: Compute Engine
            usage_families:
              - name: Compute
                resource_types:
                  - name: instances
                    operations:
                      - name: Compute
                        usage_types:
                          - n1-standard-4
                          - e2-medium

      - cloud_provider: Anthropic
        variance_profile: exponential_growth
        anomaly_profile: ""
        services:
          - name: Claude 3.5 Sonnet 2024-10-22
            usage_families:
              - name: AI Services
                resource_types: []

# Configuration Notes:
# - total_cost: Total dollar amount to generate across all providers/months
# - resource_count: Number of distinct resources per provider
# - variance_profile: linear_growth, steady_state, periodic_surge, exponential_growth
# - anomaly_profile: exponential_burst, gradual_decline, "" (none)
# - Services must match CloudZero's provider service catalogs
# - Empty resource_types: [] means generate default resource types for that service
# - Use prototype configs in configs/generate/prototype/ as detailed examples
# - Run 'czbg analyze config <file>' to validate configuration before generation
"""
    
    # Create upload.template.yaml
    upload_template_content = """# Upload Configuration - Smart Defaults Format
# Copy this file to ~/.czbg/configs/upload/ and customize for your needs

upload_config:
  # Dataset name from generated datasets (required)
  dataset_name: "REPLACE_WITH_YOUR_DATASET_NAME"

  # Global defaults (inherited by all organizations unless overridden)
  months:
    - "2025-01"
    - "2025-02"
    - "2025-03"
    - "2025-04"
    - "2025-05"
    - "2025-06"
    - "2025-07"
    - "2025-08"
    - "2025-09"
  providers:
    - Anthropic
    - AWS
    - AWS Marketplace
    - Azure
    - Confluent
    - Databricks
    - Datadog
    - Fastly
    - GCP
    - MongoDB
    - New Relic
    - OpenAI
    - Oracle Cloud
    - Redis
    - Snowflake
    - Sumologic

  # Organizations to upload to (required)
  organizations:
    DemoTest: # Used as-is for S3 folder creation
      api_token_env: "CZ_ORG_API_KEY_DEMOTEST" # Environment variable name
      # Inherits global months and providers by default
      #
      # Optional overrides:
      # months: ["2025-07", "2025-08", "2025-09"]  # Override months for this org
      # providers: ["AWS", "GCP"]                   # Override providers for this org
      # stage_overrides:                             # Optional stage overrides
      #   GCP: "usage"                              # Force GCP to use usage stage
      #   OpenAI: "metadata"                        # Force OpenAI to use metadata stage

# Configuration Notes:
# - dataset_name: Must match existing dataset from 'czbg list-datasets'
# - months: Format must be YYYY-MM (e.g., "2025-07")
# - providers: List of cloud provider names
# - Organization names are used as-is for S3 folder names
# - Environment variables must be defined in ~/.czbg/.env
# - Stage Discovery: Automatic! Stages are discovered using priority: discount > usage > metadata > k8s > generate
# - Smart Defaults: Organizations inherit global months/providers unless overridden
# - Stage Overrides: Force specific providers to use specific stages (optional)
"""
    
    # K8s template
    k8s_template_content = """# K8s Stage Configuration Template
# Copy this file to k8s.yaml and customize the dataset name

# Dataset name - should match your dataset in S3
dataset_name: "REPLACE_WITH_YOUR_DATASET_NAME"

# K8s allocation settings are constants (not configurable):
# - k8s_allocation_percentage: 0.2 (20% of eligible compute resources)
# - clusters_count: 5
# - namespaces_per_cluster: 6 (3 system + 3 custom)
# - workloads_per_namespace: 3
# - Total workloads: 5 × 6 × 3 = 90
# - Expected expansion: ~1.5x overall line items

# Configuration Notes:
# - Only eligible compute resources are allocated to K8s workloads:
#   - AWS: Service="AmazonEC2" AND ResourceType="instance"
#   - Azure: Service="microsoft.compute" AND ResourceType="virtualmachines" 
#   - GCP: Service="Compute Engine" AND ResourceType="instances"
# - Cross-month consistency is enforced (same workload assignments every month)
# - Processing validates minimum resource requirements before starting
"""

    # Metadata template
    metadata_template_content = """# Metadata Stage Configuration Template  
# Copy this file to metadata.yaml and customize for your needs

# Dataset name - should match your dataset in S3
dataset_name: "REPLACE_WITH_YOUR_DATASET_NAME"

# Majority environment for tag allocation
# Valid values: production, development, testing
majority_environment: "production"

# Regional distribution configuration
# Controls geographic distribution of resources across markets
majority_regions:
  # North America - 60% of all resources
  na:
    weight: 0.6
    majority: "us-east-1"  # Primary region within NA market
  
  # Europe, Middle East, Africa - 20% of all resources
  emea:
    weight: 0.2
    majority: "eu-west-1"  # Primary region within EMEA market
  
  # Latin America - 15% of all resources  
  latam:
    weight: 0.15
    majority: "sa-east-1"  # Primary region within LATAM market
  
  # Asia Pacific - 5% of all resources
  apac:
    weight: 0.05
    majority: "ap-southeast-1"  # Primary region within APAC market

# Base allocation profile for feature tagging
# Valid values: analytics_data, project_management, customer_service, marketing, security
base_allocation: "analytics_data"

# Configuration Notes:
# - Regional weights should sum to 1.0 (normalized if they don't)
# - Within each market: 70% to majority region, 30% to other regions
# - Environment and feature tags are created using metadata profiles from S3
# - Cross-month consistency is maintained through deterministic seeding
"""

    # Usage template
    usage_template_content = """# Usage Stage Configuration Template
# Copy this file to usage.yaml and customize the dataset name

# Dataset name - should match your dataset in S3
dataset_name: "REPLACE_WITH_YOUR_DATASET_NAME"

# Configuration Notes:
# - Usage stage transforms cost data into usage metrics with rates
# - Automatically processes all providers and months in the dataset
# - Creates realistic usage amounts and pricing rates per usage type
# - Cross-month consistency maintained through deterministic algorithms
# - Reads from metadata stage output (if available) or k8s/generate stage
"""

    # Discount template
    discount_template_content = """# Discount Stage Configuration Template
# Copy this file to discount.yaml and customize the dataset name

# Dataset name - should match your dataset in S3
dataset_name: "REPLACE_WITH_YOUR_DATASET_NAME"

# Configuration Notes:
# - Discount stage applies cloud billing discounts and savings plans
# - Processes active discount processors for each provider (AWS, Azure, GCP)
# - Creates realistic discount amounts and discount types
# - Dual-source processing: usage stage first, metadata stage for gaps
# - Cross-month consistency maintained for discount applications
# - Reads from usage stage output (preferred) or metadata stage output
"""
    
    # Write template files
    env_template_path.write_text(env_template_content)
    generate_template_path.write_text(generate_template_content)
    upload_template_path.write_text(upload_template_content)
    k8s_template_path.write_text(k8s_template_content)
    metadata_template_path.write_text(metadata_template_content)
    usage_template_path.write_text(usage_template_content)
    discount_template_path.write_text(discount_template_content)
    
    action = "Recreated" if templates_exist else "Initialized"
    typer.echo(f"✔ {action} workspace: {workspace_dir}")
    typer.echo(f"✔ Template: {env_template_path}")
    typer.echo(f"✔ Template: {generate_template_path}")
    typer.echo(f"✔ Template: {upload_template_path}")
    typer.echo(f"✔ Template: {k8s_template_path}")
    typer.echo(f"✔ Template: {metadata_template_path}")
    typer.echo(f"✔ Template: {usage_template_path}")
    typer.echo(f"✔ Template: {discount_template_path}")
    typer.echo(f"✔ Directory: {workspace_dir / 'configs' / 'generate'}/")
    typer.echo(f"✔ Directory: {workspace_dir / 'configs' / 'upload'}/")
    typer.echo(f"✔ Directory: {workspace_dir / 'configs' / 'k8s'}/")
    typer.echo(f"✔ Directory: {workspace_dir / 'configs' / 'metadata'}/")
    typer.echo(f"✔ Directory: {workspace_dir / 'configs' / 'usage'}/")
    typer.echo(f"✔ Directory: {workspace_dir / 'configs' / 'discount'}/")
    typer.echo(f"✔ Directory: {workspace_dir / 'logs'}/")
    typer.echo("")
    typer.echo("📋 Next steps:")
    typer.echo(f"  1. Copy template files:")
    typer.echo(f"     cp {env_template_path} {workspace_dir}/.env")
    typer.echo(f"     cp {upload_template_path} {workspace_dir}/configs/upload/upload.yaml")
    typer.echo(f"  2. Edit {workspace_dir}/.env with your actual API tokens")
    typer.echo(f"  3. Create generate configs in {workspace_dir}/configs/generate/")
    typer.echo(f"  4. Create upload configs in {workspace_dir}/configs/upload/")
    typer.echo("  5. Generate datasets: czbg generate run --config-path ~/.czbg/configs/generate/")
    typer.echo("  6. Upload to organizations: czbg upload --config ~/.czbg/configs/upload/upload.yaml")
    typer.echo("")
    typer.echo("🔗 Get API tokens: https://app.cloudzero.com/organization/api-keys")


def format_cost_columns(df: pl.DataFrame, cols: List[str]) -> pl.DataFrame:
    return df.with_columns(
        [
            pl.col(col)
            .round(2)
            .map_elements(lambda x: f"{x:,.2f}", return_dtype=pl.Utf8)
            .alias(f"{col}_formatted")
            for col in cols
        ]
    )


def format_count_columns(df: pl.DataFrame, cols: List[str]) -> pl.DataFrame:
    return df.with_columns(
        [
            pl.col(col)
            .map_elements(lambda x: f"{x:,.2f}", return_dtype=pl.Utf8)
            .alias(f"{col}_formatted")
            for col in cols
        ]
    )


@app.command()
def summarize(
    org: str = typer.Option(
        ..., "--org", help="Name of the org to summarize output for"
    ),
    run: Optional[str] = typer.Option(
        None, "--run", help="Timestamp of a specific run to analyze"
    ),
    provider: Optional[str] = typer.Option(
        None, "--provider", help="Filter summary to a specific cloud provider"
    ),
    notebook: bool = typer.Option(
        False,
        "--notebook",
        help="Generate Jupyter notebook template for deeper analysis",
    ),
):
    runs_dir = Path(f"cloudzero-demo-billing/orgs/{org}/billing/runs")
    if not runs_dir.exists():
        console.print(f"[red]No runs found at {runs_dir}[/red]")
        raise typer.Exit(1)

    valid_runs = sorted(runs_dir.glob("*"), key=lambda p: p.name)
    valid_run_names = [p.name for p in valid_runs if p.is_dir()]

    if run:
        if not re.match(r"^\d{4}-\d{2}-\d{2}_\d{6}$", run):
            console.print(
                "[red]Invalid --run format. Expected: YYYY-MM-DD_HHMMSS[/red]"
            )
            raise typer.Exit(1)

        if run not in valid_run_names:
            console.print(f"[red]Run folder not found: {run}[/red]")
            if valid_run_names:
                console.print(
                    "[bold]Available runs:[/bold] " + ", ".join(valid_run_names)
                )
            else:
                console.print("[yellow]No valid run folders found yet.[/yellow]")
            raise typer.Exit(1)
        run_dir = runs_dir / run
    else:
        if not valid_run_names:
            console.print(f"[red]No valid run folders found in {runs_dir}[/red]")
            raise typer.Exit(1)
        run_dir = runs_dir / valid_run_names[-1]  # Most recent by name

    # New: Look for partitioned files under generate/
    generate_dir = run_dir / "generate"
    parquet_files = list(generate_dir.rglob("*.parquet"))
    if not parquet_files:
        console.print(f"[red]No Parquet files found in {generate_dir}[/red]")
        raise typer.Exit(1)

    console.print(f"[dim]Reading {len(parquet_files)} files from {generate_dir}[/dim]")
    df = pl.concat([pl.read_parquet(pf) for pf in parquet_files])

    if provider:
        df = df.filter(pl.col("cloud_provider") == provider)
        if df.is_empty():
            console.print(
                f"[yellow]No rows found for provider '{provider}' in run {run_dir.name}[/yellow]"
            )
            raise typer.Exit(0)

    pl.Config.set_tbl_rows(100)

    console.rule("[bold cyan]Cost Overview")
    console.print(f"\n[bold green]Summarizing run:[/bold green] {run_dir.name}")
    if provider:
        console.print(f"[bold green]Filtered to provider:[/bold green] {provider}")

    console.print(f"[bold]Total Cost:[/bold] ${exploration.get_total_cost(df):,.2f}")

    console.print("\n[bold]Cost by Provider:[/bold]")
    formatted = format_cost_columns(
        exploration.get_cost_by_provider(df), cols=["total_cost"]
    )
    console.print(formatted.select(["cloud_provider", "total_cost_formatted"]))

    console.print("\n[bold]Cost by Month:[/bold]")
    formatted_monthly = format_cost_columns(
        exploration.get_cost_by_month(df), cols=["total_cost"]
    )
    console.print(
        formatted_monthly.sort("month").select(["month", "total_cost_formatted"])
    )

    console.print("\n[bold]Cost by Provider and Month:[/bold]")
    formatted_provider_month = format_cost_columns(
        exploration.get_min_max_cost_month_by_provider(df),
        cols=["min_total_cost", "max_total_cost"],
    )
    console.print(
        formatted_provider_month.sort(["cloud_provider", "min_month"]).select(
            [
                "cloud_provider",
                "min_month",
                "min_total_cost_formatted",
                "max_month",
                "max_total_cost_formatted",
            ]
        )
    )

    console.rule("[bold cyan]Line Item & Upload Estimates")
    line_items = exploration.get_min_max_line_items_month_by_provider(df)
    line_items = format_count_columns(
        line_items, cols=["min_line_items", "max_line_items"]
    )
    line_items = line_items.select(
        [
            "cloud_provider",
            "min_month",
            "min_line_items_formatted",
            "max_month",
            "max_line_items_formatted",
        ]
    )
    console.print("\n[bold]Line Items by Provider and Month:[/bold]")
    console.print(line_items)

    payloads = exploration.get_peak_payload_month_by_provider(df)
    payloads = format_count_columns(
        payloads, cols=["line_items", "estimated_mb", "estimated_payloads"]
    )
    payloads = payloads.select(
        [
            "cloud_provider",
            "month",
            "line_items_formatted",
            "estimated_mb_formatted",
            "estimated_payloads_formatted",
            "required_connections",
        ]
    )
    console.print("\n[bold]Estimated Payloads and Connections:[/bold]")
    console.print(payloads)

    console.rule("[bold cyan]Data Quality Checks")
    nulls = exploration.get_null_counts(df)
    total_nulls = sum(nulls.values())
    console.print(f"[bold]Total Nulls:[/bold] {total_nulls}")
    for col, count in nulls.items():
        if count:
            console.print(f"  - {col}: {count}")

    bad_resources = exploration.get_service_usage_resources(df)
    console.print(
        f"\n[bold]Unique 'service-usage' resource_type entries:[/bold] {bad_resources.shape[0]}"
    )

    if notebook:
        console.print(
            "\n[bold green]Notebook generation is not yet implemented.[/bold green]"
        )


@app.command("list-datasets")
def list_datasets(
    dataset: Optional[str] = typer.Option(None, "--dataset", help="Dataset name to drill down into"),
    stage: Optional[str] = typer.Option(None, "--stage", help="Stage name to drill down into"),
    provider: Optional[str] = typer.Option(None, "--provider", help="Provider name to drill down into"),
    month: Optional[str] = typer.Option(None, "--month", help="Month to show column schema for (e.g., '2024-07')"),
):
    """
    List datasets with progressive drill-down capabilities.
    
    Usage:
      czbillgen list-datasets                                    # List all datasets
      czbillgen list-datasets --dataset 2025-01-15_14-30-22     # List stages for dataset
      czbillgen list-datasets --dataset NAME --stage generate   # List providers for dataset+stage
      czbillgen list-datasets --dataset NAME --stage generate --provider AWS  # List months for dataset+stage+provider
      czbillgen list-datasets --dataset NAME --stage generate --provider AWS --month 2024-07  # Show column schema for specific month
    """
    console = Console()
    
    try:
        s3_service = S3DataService()
        
        # Level 1: List all datasets
        if not dataset:
            datasets = s3_service.list_datasets()
            if not datasets:
                console.print("[yellow]No datasets found[/yellow]")
                return
            
            console.print(f"[bold]Available datasets:[/bold]")
            for dataset_name in datasets:
                console.print(f"  • {dataset_name}")
            
            console.print(f"\n[dim]Use --dataset to drill down into a specific dataset[/dim]")
            return
        
        # Level 2: List stages for dataset
        if not stage:
            stages = s3_service.list_stages(dataset)
            if not stages:
                console.print(f"[yellow]No stages found for dataset '{dataset}'[/yellow]")
                return
            
            console.print(f"[bold]Stages for dataset '{dataset}':[/bold]")
            for stage_name in stages:
                console.print(f"  • {stage_name}")
            
            console.print(f"\n[dim]Use --stage to drill down into a specific stage[/dim]")
            return
        
        # Level 3: List providers for dataset+stage
        if not provider:
            providers = s3_service.list_providers(dataset, stage)
            if not providers:
                console.print(f"[yellow]No providers found for dataset '{dataset}' stage '{stage}'[/yellow]")
                return
            
            console.print(f"[bold]Providers for dataset '{dataset}' stage '{stage}':[/bold]")
            for provider_name in providers:
                console.print(f"  • {provider_name}")
            
            console.print(f"\n[dim]Use --provider to drill down into a specific provider[/dim]")
            return
        
        # Level 4: List months for dataset+stage+provider
        if not month:
            months_with_counts = s3_service.list_months_with_counts(dataset, stage, provider)
            if not months_with_counts:
                console.print(f"[yellow]No months found for dataset '{dataset}' stage '{stage}' provider '{provider}'[/yellow]")
                return
            
            console.print(f"[bold]Months for dataset '{dataset}' stage '{stage}' provider '{provider}':[/bold]")
            
            # Create table for months with line item counts
            table = Table(show_header=True, header_style="bold magenta")
            table.add_column("Month", style="cyan")
            table.add_column("Line Items", style="green", justify="right")
            
            for month_name, line_count in months_with_counts:
                table.add_row(month_name, f"{line_count:,}" if line_count > 0 else "Unknown")
            
            console.print(table)
            console.print(f"\n[dim]Use --month to view column schema for a specific month[/dim]")
            return
        
        # Level 5: Show column schema for specific month
        try:
            schema_info = s3_service.get_month_schema(dataset, stage, provider, month)
            
            console.print(f"[bold]Schema for dataset '{dataset}' stage '{stage}' provider '{provider}' month '{month}':[/bold]")
            console.print(f"[bold]Row count:[/bold] {schema_info['row_count']:,}")
            console.print()
            
            # Create table for column schema
            schema_table = Table(show_header=True, header_style="bold magenta")
            schema_table.add_column("Column Name", style="cyan")
            schema_table.add_column("Data Type", style="green")
            
            for column in schema_info['columns']:
                data_type = schema_info['dtypes'][column]
                schema_table.add_row(column, data_type)
            
            console.print(schema_table)
            
        except ValueError as e:
            console.print(f"[red]❌ {e}[/red]")
            return
        
    except ValueError as e:
        console.print(f"[red]❌ {e}[/red]")
        raise typer.Exit(1)
    except Exception as e:
        console.print(f"[red]❌ Failed to list datasets: {e}[/red]")
        raise typer.Exit(1)


@app.command("list-upload-datasets")
def list_upload_datasets(
    org: Optional[str] = typer.Option(None, "--org", help="Organization name to drill down into"),
    provider: Optional[str] = typer.Option(None, "--provider", help="Provider name to drill down into"),
    date_range: Optional[str] = typer.Option(None, "--date-range", help="Date range to show drops table for (e.g., '20240701-20240801')"),
    drop: Optional[str] = typer.Option(None, "--drop", help="Specific drop ID to show detailed file information for"),
):
    """
    List upload-ready datasets with progressive drill-down capabilities.
    
    Usage:
      czbillgen list-upload-datasets                                          # List all organizations
      czbillgen list-upload-datasets --org DemoTest                          # List providers for organization  
      czbillgen list-upload-datasets --org DemoTest --provider AWS           # List date ranges for org+provider
      czbillgen list-upload-datasets --org DemoTest --provider AWS --date-range 20240701-20240801  # Show drops table
      czbillgen list-upload-datasets --org DemoTest --provider AWS --date-range 20240701-20240801 --drop 20250722T210122Z  # Show drop details
    """
    console = Console()
    
    try:
        s3_service = S3DataService()
        
        # Level 1: List all organizations
        if not org:
            organizations = s3_service.list_upload_organizations()
            if not organizations:
                console.print("[yellow]No organizations found in upload-ready datasets[/yellow]")
                return
            
            console.print(f"[bold]Available organizations in upload-ready datasets:[/bold]")
            for org_name in organizations:
                console.print(f"  • {org_name}")
            
            console.print(f"\n[dim]Use --org to drill down into a specific organization[/dim]")
            return
        
        # Level 2: List providers for organization
        if not provider:
            providers = s3_service.list_upload_providers(org)
            if not providers:
                console.print(f"[yellow]No providers found for organization '{org}' in upload-ready datasets[/yellow]")
                return
            
            console.print(f"[bold]Providers for organization '{org}' in upload-ready datasets:[/bold]")
            for provider_name in providers:
                console.print(f"  • {provider_name}")
            
            console.print(f"\n[dim]Use --provider to drill down into a specific provider[/dim]")
            return
        
        # Level 3: List date ranges for organization + provider
        if not date_range:
            date_ranges = s3_service.list_upload_date_ranges(org, provider)
            if not date_ranges:
                console.print(f"[yellow]No date ranges found for organization '{org}' provider '{provider}' in upload-ready datasets[/yellow]")
                return
            
            console.print(f"[bold]Date ranges for organization '{org}' provider '{provider}' in upload-ready datasets:[/bold]")
            
            # Create table for date ranges
            table = Table(show_header=True, header_style="bold magenta")
            table.add_column("Date Range", style="cyan")
            table.add_column("Period", style="green")
            
            for date_range_name in date_ranges:
                # Parse date range to show readable period
                try:
                    start_date = date_range_name.split('-')[0]
                    end_date = date_range_name.split('-')[1]
                    # Format as YYYY-MM-DD for readability
                    start_formatted = f"{start_date[:4]}-{start_date[4:6]}-{start_date[6:8]}"
                    end_formatted = f"{end_date[:4]}-{end_date[4:6]}-{end_date[6:8]}"
                    period = f"{start_formatted} to {end_formatted}"
                except:
                    period = "Invalid format"
                
                table.add_row(date_range_name, period)
            
            console.print(table)
            console.print(f"\n[dim]Use --date-range to view manifest and data info for a specific date range[/dim]")
            return
        
        # Level 4: Show drops table for specific date range
        if not drop:
            try:
                drops_info = s3_service.list_drops_in_date_range(org, provider, date_range)
                
                if not drops_info:
                    console.print(f"[yellow]No drops found for organization '{org}' provider '{provider}' date range '{date_range}'[/yellow]")
                    return
                
                console.print(f"[bold]Drops for organization '{org}' provider '{provider}' date range '{date_range}':[/bold]")
                
                # Create table for drops
                table = Table(show_header=True, header_style="bold magenta")
                table.add_column("Drop ID", style="cyan")
                table.add_column("Active", style="green", justify="center")
                table.add_column("Files", style="yellow", justify="right")
                table.add_column("Total Rows", style="blue", justify="right")
                
                # Get active drop ID from manifest
                active_drop_id = drops_info[0].get("active_drop_id") if drops_info else None
                
                for drop_info in drops_info:
                    drop_id = drop_info["drop_id"]
                    is_active = "✓" if drop_id == active_drop_id else ""
                    file_count = drop_info["file_count"]
                    total_rows = drop_info["total_rows"]
                    
                    table.add_row(
                        drop_id,
                        is_active,
                        f"{file_count:,}" if file_count > 0 else "0",
                        f"{total_rows:,}" if total_rows > 0 else "0"
                    )
                
                console.print(table)
                console.print(f"\n[dim]Use --drop to view detailed file information for a specific drop[/dim]")
                return
                
            except ValueError as e:
                console.print(f"[red]❌ {e}[/red]")
                return
        
        # Level 5: Show detailed drop information
        try:
            drop_details = s3_service.get_drop_details(org, provider, date_range, drop)
            
            console.print(f"[bold]Drop details for '{drop}' in organization '{org}' provider '{provider}' date range '{date_range}':[/bold]")
            
            # Display drop summary
            console.print(f"[bold]Drop Information:[/bold]")
            console.print(f"  • Drop ID: {drop_details['drop_id']}")
            console.print(f"  • Active Drop: {'Yes' if drop_details.get('is_active', False) else 'No'}")
            console.print(f"  • File Count: {drop_details['file_count']:,}")
            console.print(f"  • Total Rows: {drop_details['total_rows']:,}")
            console.print()
            
            # Display file information
            if drop_details['files']:
                console.print(f"[bold]Data Files:[/bold]")
                
                # Create table for files
                files_table = Table(show_header=True, header_style="bold magenta")
                files_table.add_column("File Name", style="cyan")
                files_table.add_column("Size (MB)", style="yellow", justify="right")
                files_table.add_column("Rows", style="blue", justify="right")
                files_table.add_column("Last Modified", style="green")
                
                for file_info in drop_details['files']:
                    file_size_mb = file_info['size_bytes'] / (1024 * 1024)
                    files_table.add_row(
                        file_info['filename'],
                        f"{file_size_mb:.2f}",
                        f"{file_info['rows']:,}" if file_info['rows'] > 0 else "Unknown",
                        file_info['last_modified']
                    )
                
                console.print(files_table)
                console.print()
            
            # Display column schema if available
            if drop_details.get('columns'):
                console.print(f"[bold]Column Schema:[/bold]")
                
                # Create table for columns
                columns_table = Table(show_header=True, header_style="bold magenta")
                columns_table.add_column("Column Name", style="cyan")
                
                for column in drop_details['columns']:
                    columns_table.add_row(column)
                
                console.print(columns_table)
            
            # Display schema error if any
            if drop_details.get('schema_error'):
                console.print(f"\n[yellow]⚠️  Schema read warning: {drop_details['schema_error']}[/yellow]")
            
        except ValueError as e:
            console.print(f"[red]❌ {e}[/red]")
            return
        
    except ValueError as e:
        console.print(f"[red]❌ {e}[/red]")
        raise typer.Exit(1)
    except Exception as e:
        console.print(f"[red]❌ Failed to list upload datasets: {e}[/red]")
        raise typer.Exit(1)


if __name__ == "__main__":
    app()
