import json
import pathlib
import click
from prettytable import PrettyTable
from sigma.processing.resolver import SigmaPipelineNotFoundError

from sigma.cli.convert import pipeline_resolver
from sigma.cli.rules import check_rule_errors, load_rules
from sigma.analyze.attack import score_functions, calculate_attack_scores
from sigma.analyze.fields import extract_fields_from_collection
from sigma.analyze.stats import create_logsourcestats, format_row
from sigma.rule import SigmaLevel, SigmaStatus
from sigma.plugins import InstalledSigmaPlugins
from sigma.conversion.base import Backend


@click.group(name="analyze", help="Analyze Sigma rule sets")
def analyze_group():
    pass


@analyze_group.command(
    name="attack",
    help="Create MITRE™️ ATT&CK heatmaps from Sigma rule set. Score functions are: "
    + ", ".join(
        (f"{definition[1]} ({func})" for func, definition in score_functions.items())
    ),
)
@click.option(
    "--file-pattern",
    "-P",
    default="*.yml",
    show_default=True,
    help="Pattern for file names to be included in recursion into directories.",
)
@click.option(
    "--min-level",
    "-L",
    default="INFORMATIONAL",
    help="The minimun level of the rule to be include.",
)
@click.option(
    "--min-status",
    "-T",
    default="UNSUPPORTED",
    help="The minimun status of the rule to be include.",
)
@click.option(
    "--subtechniques/--no-subtechniques",
    "-s/-S",
    default=True,
)
@click.option(
    "--max-color",
    "-c",
    default="#ff0000",
    show_default=True,
    help="Color used for maximum score.",
)
@click.option(
    "--min-color",
    "-C",
    default="#ffffff00",
    show_default=True,
    help="Color used for zero score.",
)
@click.option(
    "--max-score",
    "-m",
    type=int,
    default=None,
    help="Set fixed maximum score. All scores above are rendered as maximum. Increases color scale resolution for scores below.",
)
@click.option(
    "--min-score",
    "-M",
    type=int,
    default="0",
    show_default=True,
    help="Minimum score. All scores below are not explicitly colored.",
)
@click.argument(
    "function",
    type=click.Choice(score_functions.keys()),
)
@click.argument(
    "output",
    type=click.File("w"),
)
@click.argument(
    "input",
    nargs=-1,
    required=True,
    type=click.Path(exists=True, allow_dash=True, path_type=pathlib.Path),
)
def analyze_attack(
    file_pattern,
    min_level,
    min_status,
    subtechniques,
    max_color,
    min_color,
    max_score,
    min_score,
    function,
    output,
    input,
):
    try:
        min_sigmalevel = SigmaLevel[min_level.upper()]
    except:
        min_sigmalevel = SigmaLevel.INFORMATIONAL
    try:
        min_sigmastatus = SigmaStatus[min_status.upper()]
    except:
        min_sigmastatus = SigmaStatus.UNSUPPORTED
    from sigma.data.mitre_attack import (
        mitre_attack_techniques_tactics_mapping,
        mitre_attack_version,
    )

    rules = load_rules(input, file_pattern)
    check_rule_errors(rules)
    score_function = score_functions[function][0]
    scores = calculate_attack_scores(rules, score_function, not subtechniques,min_sigmalevel=min_sigmalevel,min_sigmastatus=min_sigmastatus,)
    layer_techniques = [
        {
            "techniqueID": technique,
            "tactic": tactic,
            "score": score,
            "color": "",
            "comment": "",
            "enabled": True,
            "metadata": [],
            "links": [],
            "showSubtechniques": False,
        }
        for technique, score in scores.items()
        for tactic in mitre_attack_techniques_tactics_mapping.get(technique, [])
    ]
    layer = {
        "name": "Sigma Analytics Coverage",
        "versions": {
            "attack": mitre_attack_version,
            "navigator": "4.8.1",
            "layer": "4.4",
        },
        "domain": "enterprise-attack",
        "description": f"Sigma coverage heatmap generated by Sigma CLI with score function {function}",
        "gradient": {
            "colors": [
                min_color,
                max_color,
            ],
            "minValue": min_score,
            "maxValue": max_score or max(scores.values()),
        },
        "techniques": layer_techniques,
    }
    json.dump(layer, output, indent=2)

@analyze_group.command(name="logsource", help="Create stats about logsources.")
@click.option(
    "--file-pattern",
    "-P",
    default="*.yml",
    show_default=True,
    help="Pattern for file names to be included in recursion into directories.",
)
@click.option(
    "--sort-by",
    "-k",
    type=str,
    default="Overall",
    show_default=True,
    help="Sort by column.",
)
@click.argument(
    "output",
    type=click.File("w"),
)
@click.argument(
    "input",
    nargs=-1,
    required=True,
    type=click.Path(exists=True, allow_dash=True, path_type=pathlib.Path),
)
def analyze_logsource(
    file_pattern,
    sort_by,
    output,
    input,
):
    rules = load_rules(input, file_pattern)
    check_rule_errors(rules)
    stats = create_logsourcestats(rules)

    # Extract column header
    headers = ["Logsource"] + list(next(iter(stats.values())).keys())

    # Prepare rows
    rows = [[key] + list(value.values()) for key, value in stats.items()]
    sort_index = headers.index(sort_by)
    rows.sort(key=lambda x: x[sort_index], reverse=True)

    # Determine col width
    column_widths = [
        max(len(str(item)) for item in column) for column in zip(*([headers] + rows))
    ]

    # Print table
    print("-+-".join("-" * width for width in column_widths), file=output)
    print(format_row(headers, column_widths), file=output)
    print("-+-".join("-" * width for width in column_widths), file=output)
    for row in rows:
        print(format_row(row, column_widths), file=output)


@analyze_group.command(
    name="fields",
    help="Extract field names from Sigma rules for a given target backend and processing pipeline(s).",
)
@click.option(
    "--file-pattern",
    "-P",
    default="*.yml",
    show_default=True,
    help="Pattern for file names to be included in recursion into directories.",
)
@click.option(
    "--target",
    "-t",
    type=str,
    required=True,
    help="Target backend to use for field name escaping and quoting.",
)
@click.option(
    "--pipeline",
    "-p",
    multiple=True,
    help="Specify processing pipelines as identifiers ("
    + click.style("sigma list pipelines", bold=True, fg="green")
    + ") or YAML files or directories",
)
@click.option(
    "--pipeline-check/--disable-pipeline-check",
    default=True,
    help="Verify if a pipeline is used that is intended for another backend.",
)
@click.option(
    "--group/--no-group",
    default=False,
    help="Group fields by logsource.",
)
@click.option(
    "--enable-template-vars",
    is_flag=True,
    default=False,
    help="Enable template variable support in processing pipeline. WARNING: This feature can be dangerous and allow arbitrary code execution if used with untrusted Sigma rules.",
)
@click.option(
    "--template-vars-path",
    multiple=True,
    type=click.Path(exists=True, path_type=pathlib.Path),
    help="Allowed paths for template variable expansion. Can be specified multiple times.",
)
@click.argument(
    "input",
    nargs=-1,
    required=True,
    type=click.Path(exists=True, allow_dash=True, path_type=pathlib.Path),
)
def analyze_fields(file_pattern, target, pipeline, pipeline_check, group, enable_template_vars, template_vars_path, input):
    """Extract field names from Sigma rule sets.
    
    This command extracts and outputs all unique field names present in the given
    Sigma rule collection, formatted for the specified target backend.
    """
    # Load plugins and get available backends
    plugins = InstalledSigmaPlugins.autodiscover()
    backends = plugins.backends
    
    if target not in backends:
        available_targets = ", ".join(sorted(backends.keys()))
        raise click.ClickException(
            f"Unknown target '{target}'. Available targets are: {available_targets}"
        )
    
    # Load rules
    rules = load_rules(input, file_pattern)
    check_rule_errors(rules)
    
    # Resolve pipelines
    try:
        processing_pipeline = pipeline_resolver.resolve(
            pipeline, target if pipeline_check else None
        )
        
        # Configure template variable settings on the processing pipeline
        if enable_template_vars:
            processing_pipeline.allow_template_vars = True
        if template_vars_path:
            processing_pipeline.vars_allowed_paths = [str(p) for p in template_vars_path]
    except SigmaPipelineNotFoundError as e:
        raise click.UsageError(
            f"The pipeline '{e.spec}' was not found.\n"
            + "List all installed processing pipelines with: "
            + click.style(f"sigma list pipelines {target}", bold=True, fg="green")
            + "\n"
            "List pipeline plugins for installation with: "
            + click.style(
                f"sigma plugin list --plugin-type pipeline", bold=True, fg="green"
            )
            + "\n"
            + "Pipelines not listed here are treated as file names."
        )
    
    # Initialize backend
    backend_class = backends[target]
    try:
        backend: Backend = backend_class(
            processing_pipeline=processing_pipeline,
            collect_errors=True,
        )
    except Exception as e:
        raise click.ClickException(f"Failed to initialize backend '{target}': {str(e)}")
    
    # Extract fields
    all_fields, errors = extract_fields_from_collection(rules, backend, group)
    
    # Handle errors
    if errors:
        click.echo("Warnings during field extraction:", err=True)
        for error in errors:
            click.echo(f"* {error}", err=True)

    if group:
        table = PrettyTable()
        table.field_names = ["Logsource", "Fields"]
        table.align["Logsource"] = "r"
        table.align["Fields"] = "l"
        for logsource, fields in sorted(all_fields.items()):
            table.add_row([logsource, "\n".join(sorted(fields))])
        click.echo(table)
    else:
        # Output fields sorted
        click.echo("\n".join(sorted(all_fields)))
