"""
Query Commands - Typer CLI for SQL-based dataset queries.

Commands:
- less: Interactive dataset browser with pagination
- sites: Site-specific queries with URL filtering
- analyze: Transaction log analysis
- aggregate: Complex data aggregation
- slice: Dataset slicing/subsetting
- stream: Real-time data streaming
- warc: WARC file extraction and management
"""
from typing import Optional
import typer

from owilix.cli._common.context import CLIContext
from owilix.cli._common.executor import execute_command
from owilix.core.tasks.query import (
    query_less, query_sites, query_analyze, query_aggregate,
    slice_query, stream_query, query_stats
)
from owilix.core.tasks.warc.query_warc import warc as query_warc
from owilix.cli._common.output import OutputWriter

# Create sub-app for query commands
app = typer.Typer(
    name="query",
    help="SQL-based dataset query operations",
    no_args_is_help=True,
)


@app.command()
def less(
    ctx: typer.Context,
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier (e.g., 'main:latest')"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier (e.g., 'lrz:latest')"),
    select: str = typer.Option("url,title,main_content", "--select", "-s", help="SQL SELECT columns"),
    where: str = typer.Option("", "--where", "-w", help="SQL WHERE clause"),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Maximum rows to return"),
    files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"),
    page_size: int = typer.Option(10, "--page-size", help="Results per page"),
    batch_size: int = typer.Option(100, "--batch-size", help="Rows per processing batch"),
    pq_batch_size: int = typer.Option(10, "--pq-batch", help="Parquet files per batch"),
    verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"),
    job_name: Optional[str] = typer.Option(None, "--job", help="Job name for transaction logging"),
    resume: bool = typer.Option(False, "--resume", help="Resume from previous execution"),
    async_mode: bool = typer.Option(False, "--async/--sync", help="Use async executor (sync streams results, async collects all first)"),
    prefetch: int = typer.Option(3, "--prefetch", help="Number of parallel query batches"),
    search: Optional[str] = typer.Option(None, "--search", "-S", help="Full-text search terms (ANDed with --where)"),
    search_language: str = typer.Option("eng", "--search-lang", help="Search language (ISO 639-3, e.g. eng, deu)"),
    search_all: bool = typer.Option(False, "--search-all", help="Require ALL search terms (default: any)"),
    search_limit: int = typer.Option(100, "--search-limit", help="Max documents from remote index search"),
):
    """
    Interactive dataset browser with SQL query capabilities.
    
    Browse dataset content with pagination, filtering, and SQL-based selection.
    Results can be displayed in various formats using the global --format flag.
    
    Examples:
        owi query less -R lrz:latest
        owi --format json query less -R lrz:latest --limit 100
        owi --format jsonl -o results.jsonl query less -L main:latest
        owi --yes query less -R lrz:latest  # Auto-confirm pagination
    """
    cli_ctx: CLIContext = ctx.obj
    
    result = execute_command(
        cli_ctx,
        query_less,
        manager=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        select=select,
        where=where,
        limit=limit,
        files=files,
        format=cli_ctx.output_format,
        output_file=cli_ctx.output_file,
        auto_confirm=cli_ctx.auto_yes,
        page_size=page_size,
        batch_size=batch_size,
        pq_batch_size=pq_batch_size,
        console=cli_ctx.console,
        verbose=verbose,
        job_name=job_name,
        resume=resume,
        async_mode=async_mode,
        prefetch=prefetch,
        search_terms=search,
        search_language=search_language,
        search_conjunctive=search_all,
        search_limit=search_limit,
        command_name="query less"
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)


@app.command()
def sites(
    ctx: typer.Context,
    urls_file: str = typer.Argument(..., help="File containing URLs for filtering"),
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier"),
    select: str = typer.Option("url,url_subdomain,url_domain,url_suffix", "--select", "-s", help="SQL SELECT columns"),
    where: str = typer.Option("", "--where", "-w", help="Additional WHERE conditions"),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Maximum rows to return"),
    files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"),
    page_size: int = typer.Option(10, "--page-size", help="Results per page"),
    batch_size: int = typer.Option(100, "--batch-size", help="Rows per processing batch"),
    pq_batch_size: int = typer.Option(1, "--pq-batch", help="Parquet files per batch"),
    verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"),
    job_name: Optional[str] = typer.Option(None, "--job", help="Job name for transaction logging"),
    resume: bool = typer.Option(False, "--resume", help="Resume from previous execution"),
    async_mode: bool = typer.Option(False, "--async/--sync", help="Use async executor"),
    prefetch: int = typer.Option(3, "--prefetch", help="Number of parallel query batches"),
):
    """
    Execute site-specific queries using URL filtering.
    
    Retrieves site-specific information based on URLs in a file.
    Supports domain, subdomain, and path filtering with wildcard support.
    Use global --format flag for output format (table, json, jsonl).
    
    Examples:
        owi query sites urls.txt -L main:latest
        owi --format json query sites urls.txt -R lrz:latest
    """
    cli_ctx: CLIContext = ctx.obj
    
    result = execute_command(
        cli_ctx,
        query_sites,
        manager=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        urls_file=urls_file,
        select=select,
        where=where if where else None,
        limit=limit,
        files=files,
        format=cli_ctx.output_format,
        output_file=cli_ctx.output_file,
        auto_confirm=cli_ctx.auto_yes,
        page_size=page_size,
        batch_size=batch_size,
        pq_batch_size=pq_batch_size,
        console=cli_ctx.console,
        verbose=verbose,
        job_name=job_name,
        resume=resume,
        async_mode=async_mode,
        prefetch=prefetch,
        command_name="query sites"
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)


@app.command()
def analyze(
    ctx: typer.Context,
    job_name: str = typer.Argument(..., help="Job name to analyze"),
    details: bool = typer.Option(False, "--details", "-d", help="Show detailed information"),
    status_filter: Optional[str] = typer.Option(None, "--status", help="Filter by status"),
    error_limit: int = typer.Option(10, "--error-limit", help="Max errors to show"),
):
    """
    Analyze transaction log for a specific job.
    """
    cli_ctx: CLIContext = ctx.obj
    
    result = execute_command(
        cli_ctx,
        query_analyze,
        manager=cli_ctx.owi,
        job_name=job_name,
        details=details,
        status_filter=status_filter,
        error_limit=error_limit,
        console=cli_ctx.console,
        command_name="query analyze"
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)


@app.command()
def aggregate(
    ctx: typer.Context,
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier"),
    select: str = typer.Option("url_suffix,count(url) as c", "--select", "-s", help="SQL SELECT with aggregation"),
    where: str = typer.Option("", "--where", "-w", help="SQL WHERE clause"),
    groupby: str = typer.Option("url_suffix", "--groupby", "-g", help="SQL GROUP BY clause"),
    aggregate_query: str = typer.Option(
        "SELECT * FROM aggregates ORDER BY c DESC LIMIT 20",
        "--aggregate", "-a",
        help="Final aggregation SQL query"
    ),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Per-file row limit"),
    files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"),
    batch_size: int = typer.Option(100, "--batch-size", help="Rows per batch"),
    pq_batch_size: int = typer.Option(1, "--pq-batch", help="Parquet files per batch"),
    verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"),
    async_mode: bool = typer.Option(False, "--async/--sync", help="Use async executor"),
    prefetch: int = typer.Option(3, "--prefetch", help="Number of parallel query batches"),
):
    """
    Run SQL aggregation queries across datasets.
    """
    cli_ctx: CLIContext = ctx.obj
    
    result = execute_command(
        cli_ctx,
        query_aggregate,
        manager=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        select=select,
        where=where if where else None,
        groupby=groupby,
        aggregate_query=aggregate_query,
        limit=limit,
        files=files,
        format=cli_ctx.output_format,
        output_file=cli_ctx.output_file,
        batch_size=batch_size,
        pq_batch_size=pq_batch_size,
        async_mode=async_mode,
        prefetch=prefetch,
        console=cli_ctx.console,
        command_name="query aggregate"
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)


@app.command()
def stats(
    ctx: typer.Context,
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier"),
    files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"),
    topk: int = typer.Option(20, "--topk", "-k", help="Number of top items per category"),
    write_back: bool = typer.Option(False, "--write-back", "-w", help="Save stats to dataset .stats folder"),
    force_write: bool = typer.Option(False, "--force", help="Force overwrite existing stats"),
    async_mode: bool = typer.Option(False, "--async/--sync", help="Use async executor"),
    prefetch: int = typer.Option(3, "--prefetch", help="Number of parallel query batches"),
):
    """
    Calculate statistics over datasets.
    
    Statistics include record counts, URL distributions, date ranges, and MIME types.
    Uses global --format flag for output (table, json, jsonl).
    
    Examples:
        owi query stats -R lexis:latest/collectionName=cefal
        owi --format json query stats -R lexis:latest -o stats.json
        owi query stats -R lexis:latest --topk 50
    """
    cli_ctx: CLIContext = ctx.obj
    
    result = execute_command(
        cli_ctx,
        query_stats,
        manager=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        files=files,
        topk=topk,
        format=cli_ctx.output_format,
        output_file=cli_ctx.output_file,
        write_back=write_back,
        force_write=force_write,
        async_mode=async_mode,
        prefetch=prefetch,
        console=cli_ctx.console,
        command_name="query stats"
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)


@app.command()
def slice(
    ctx: typer.Context,
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier"),
    select: str = typer.Option("*", "--select", "-s", help="SQL SELECT columns"),
    where: str = typer.Option("", "--where", "-w", help="SQL WHERE clause"),
    limit: str = typer.Option("", "--limit", "-l", help="Row limit"),
    paths: str = typer.Option("**/", "--paths", "-p", help="Path glob pattern"),
    creator: Optional[str] = typer.Option(None, "--creator", help="Creator identifier"),
    resource_type: str = typer.Option("owip", "--type", "-t", help="Resource type"),
    access: str = typer.Option("public", "--access", "-a", help="Access level"),
    collection_name: str = typer.Option("userslice", "--collection", "-c", help="Collection name"),
    urls_file: Optional[str] = typer.Option(None, "--urls", help="File with URLs to filter"),
    internal_id: Optional[str] = typer.Option(None, "--id", help="Existing dataset ID to update"),
    partitioned_by: str = typer.Option("", "--partition", help="Partition columns"),
    overwrite: bool = typer.Option(True, "--overwrite/--no-overwrite", help="Overwrite existing"),
    slice_ciff: bool = typer.Option(True, "--ciff/--no-ciff", help="Also slice CIFF files"),
    num_threads: int = typer.Option(4, "--threads", help="Threads for CIFF processing"),
    batch_size: int = typer.Option(500, "--batch-size", help="Rows per batch"),
    pq_batch_size: int = typer.Option(1, "--pq-batch", help="Parquet files per batch"),
    chunk_size: int = typer.Option(1000000, "--chunk-size", help="Rows per chunk"),
    job_name: Optional[str] = typer.Option(None, "--job", help="Job name for logging"),
    resume: bool = typer.Option(False, "--resume", help="Resume previous execution"),
    verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"),
    yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation"),
    # Output format options using global enums from main
    format: Optional[str] = typer.Option(None, "--format", help="Output format (table, json, jsonl, json+, jsonl+)"),
    output_file: Optional[str] = typer.Option(None, "--output", "-o", help="File to write output to"),
    prefetch: int = typer.Option(3, "--prefetch", help="Number of parallel query batches"),
    search: Optional[str] = typer.Option(None, "--search", "-S", help="Full-text search terms (ANDed with --where)"),
    search_language: str = typer.Option("eng", "--search-lang", help="Search language (ISO 639-3, e.g. eng, deu)"),
    search_all: bool = typer.Option(False, "--search-all", help="Require ALL search terms (default: any)"),
    search_limit: int = typer.Option(100, "--search-limit", help="Max documents from remote index search"),
):
    """
    Create a new dataset by slicing data from existing datasets.
    """
    cli_ctx: CLIContext = ctx.obj
    # Resolve format: command-line arg > global arg > default
    format = format or cli_ctx.output_format or "table"
    
    # Perform slice directly (as it prints progress etc extensively)
    result = slice_query(
        owi=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        select=select,
        where=where,
        limit=limit,
        paths=paths,
        creator=creator,
        resource_type=resource_type,
        access=access,
        collection_name=collection_name,
        urls_file=urls_file,
        internal_id=internal_id,
        partitioned_by=partitioned_by,
        overwrite_ignore=overwrite,
        slice_ciff=slice_ciff,
        num_threads=num_threads,
        batch_size=batch_size,
        pq_batch_size=pq_batch_size,
        chunk_size=chunk_size,
        job_name=job_name,
        resume=resume,
        verbose=verbose,
        console=cli_ctx.console if format == "table" and not output_file else None,
        autoyes=yes,
        prefetch=prefetch,
        search_terms=search,
        search_language=search_language,
        search_conjunctive=search_all,
        search_limit=search_limit,
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)
    
    if result and result.success:
        if format in ("json", "jsonl", "json+", "jsonl+") or output_file:
             # Structured output for automation
             output_data = result.object
             if not output_data:
                 output_data = {"status": "success", "message": "Slice completed successfully"}
             
             # Update context for OutputWriter
             cli_ctx.output_format = format
             cli_ctx.output_file = output_file
             
             with OutputWriter(cli_ctx) as writer:
                 writer.write_records([output_data])
        else:
             cli_ctx.console.print(f"[green]✓ Slice completed successfully[/green]")


@app.command()
def stream(
    ctx: typer.Context,
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier"),
    select: str = typer.Option("url,title,main_content", "--select", "-s", help="SQL SELECT columns"),
    where: str = typer.Option("", "--where", "-w", help="SQL WHERE clause"),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Row limit"),
    files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"),
    consumer: str = typer.Option(
        "owilix.core.stream.ConsumeToStdouts",
        "--consumer",
        help="Consumer class for processing batches"
    ),
    batch_size: int = typer.Option(100, "--batch-size", help="Rows per batch"),
    pq_batch_size: int = typer.Option(1, "--pq-batch", help="Parquet files per batch"),
    prefetch: int = typer.Option(3, "--prefetch", help="Batches to prefetch"),
    queue_size: int = typer.Option(5, "--queue", help="Buffer queue size"),
    verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"),
    format: Optional[str] = typer.Option(None, "--format", help="Output format (table, json, jsonl, json+, jsonl+)"),
    output_file: Optional[str] = typer.Option(None, "--output", "-o", help="File to write output to"),
):
    """
    Stream query results to stdout or custom consumer.
    """
    cli_ctx: CLIContext = ctx.obj
    # Resolve format: command-line arg > global arg > default
    format = format or cli_ctx.output_format or "table"
    
    result = stream_query(
        owi=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        select=select,
        where=where,
        limit=limit,
        files=files,
        consumer=consumer,
        batch_size=batch_size,
        pq_batch_size=pq_batch_size,
        prefetch=prefetch,
        queue_size=queue_size,
        verbose=verbose,
        console=cli_ctx.console,
    )
    
    if result and not result.success:
        raise typer.Exit(code=1)


@app.command()
def warc(
    ctx: typer.Context,
    local: str = typer.Option("", "--local", "-L", help="Local dataset specifier"),
    remote: str = typer.Option("", "--remote", "-R", help="Remote dataset specifier"),
    where: str = typer.Option("", "--where", "-w", help="SQL WHERE clause"),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Row limit"),
    files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"),
    output_dir: Optional[str] = typer.Option(None, "--output", "-o", help="Output directory or log path"),
    warc_config: Optional[str] = typer.Option(None, "--warc-config", help="WARC location config file"),
    warc_token: Optional[str] = typer.Option(None, "--warc-token", help="Access token for pre-signed URLs (overrides OWI_WARC_ACCESS_TOKEN env var)"),
    use_presign: bool = typer.Option(True, "--use-presign/--no-presign", help="Use pre-signed URLs for WARC downloads (default: True)"),
    batch_size: int = typer.Option(100, "--batch-size", help="Rows per batch"),
    pq_batch_size: int = typer.Option(1, "--pq-batch", help="Parquet files per batch"),
    verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"),
    job_name: Optional[str] = typer.Option(None, "--job", help="Job name for logging"),
):
    """
    Extract WARC file locations from datasets and download WARC records.

    Pre-signed URL authentication (when use_presign=True):
        Token is loaded from (in priority order):
        1. --warc-token parameter
        2. OWI_WARC_ACCESS_TOKEN environment variable
        3. ~/.s3_access_token file

    To use .env-rc file for token:
        source .env-rc
        owi query warc --remote "lrz:latest" --warc-config .env-warc-cfg.json

    To disable pre-signed URLs and use direct S3 credentials:
        owi query warc --no-presign --remote "lrz:latest" --warc-config .env-warc-cfg.json
    """
    cli_ctx: CLIContext = ctx.obj

    result = execute_command(
        cli_ctx,
        query_warc,
        manager=cli_ctx.owi,
        local_specifier=local,
        remote_specifier=remote,
        where=where,
        limit=limit,
        files=files,
        log_path=output_dir if output_dir else "",
        warc_location_cfg=warc_config,
        warc_token=warc_token,
        use_presign=use_presign,
        batch_size=batch_size,
        pq_batch_size=pq_batch_size,
        console=cli_ctx.console,
        verbose=verbose,
        group_name=job_name if job_name else "",
        command_name="query warc"
    )

    if result and not result.success:
        raise typer.Exit(code=1)