from __future__ import annotations
import os
import json
from typing import Literal, Sequence
import click
import fsspec
from yarl import URL
from matchviz import (
    create_neuroglancer_state,
    fetch_all_matches,
    get_tilegroup_s3_url,
    read_bigstitcher_xml,
    save_interest_points,
    summarize_matches,
)
from matchviz.neuroglancer_styles import (
    NeuroglancerViewerStyle,
    neuroglancer_view_styles,
)
import structlog

@click.group("matchviz")
def cli(): ...

log_level = click.option('--log-level', type=click.STRING, default='info')

@cli.command("save-points")
@click.option("--src", type=click.STRING, required=True)
@click.option("--dest", type=click.STRING, required=True)
def save_interest_points_cli(src: str, dest: str):
    """
    Save bigstitcher interest points from n5 to neuroglancer precomputed annotations.
    """
    # strip trailing '/' from src and dest
    src_parsed = src.rstrip("/")
    dest_parsed = dest.rstrip("/")
    save_points(url=src_parsed, dest=dest_parsed)


def save_points(url: str, dest: str):
    bs_model = read_bigstitcher_xml(os.path.join(url, "bigstitcher.xml"))
    save_interest_points(bs_model=bs_model, base_url=url, out_prefix=dest)


@cli.command("ngjson")
@click.option("--alignment-url", type=click.STRING, required=True)
@click.option("--points-url", type=click.STRING)
@click.option("--matches-url", type=click.STRING)
@click.option("--dest-path", type=click.STRING, required=True)
@click.option("--style", type=click.STRING, multiple=True)
def save_neuroglancer_json_cli(
    alignment_url: str,
    dest_path: str,
    points_url: str | None,
    matches_url: str | None,
    style: Sequence[NeuroglancerViewerStyle] | None = None,
):
    """
    Generate a neuroglancer viewer state as a JSON document.
    """
    log = structlog.get_logger(__name__)
    alignment_url_parsed = URL(alignment_url)
    if points_url is not None:
        points_url_parsed = URL(points_url)
    else:
        points_url_parsed = None

    if matches_url is not None:
        matches_url_parsed = URL(matches_url)
    else:
        matches_url_parsed = None

    dest_path_parsed = dest_path.rstrip("/")
    if style is None or len(style) < 1:
        style = neuroglancer_view_styles
    for _style in style:
        out_path = save_neuroglancer_json(
            alignment_url=alignment_url_parsed,
            dest_path=dest_path_parsed,
            points_url=points_url_parsed,
            matches_url=matches_url_parsed,
            style=_style,
        )
        log.info(f"Saved neuroglancer JSON state for style {_style} to {out_path}")


def save_neuroglancer_json(
    *,
    alignment_url: str,
    points_url: str | None,
    matches_url: str | None,
    dest_path: str,
    style: NeuroglancerViewerStyle,
) -> str:
    bs_model = read_bigstitcher_xml(alignment_url.joinpath("bigstitcher.xml"))
    tilegroup_s3_url = get_tilegroup_s3_url(bs_model)
    state = create_neuroglancer_state(
        image_url=tilegroup_s3_url, points_url=points_url, matches_url = matches_url, style=style
    )
    out_fname = f"{style}.json"
    out_path = os.path.join(dest_path, out_fname)
    if dest_path.startswith("s3://"):
        fs, _ = fsspec.url_to_fs(dest_path)
    else:
        fs, _ = fsspec.url_to_fs(dest_path, auto_mkdir=True)

    with fs.open(out_path, mode="w") as fh:
        fh.write(json.dumps(state.to_json(), indent=2))

    return out_path

@cli.command('tabulate-matches')
@click.option('--alignment-url', type=click.STRING, required=True)
@click.option('--output', type=click.STRING, default='csv')
def tabulate_matches_cli(alignment_url: str, output: Literal['csv'] | None):
    """
    Generate a tabular representation of the correspondence metadata generated by bigstitcher. 
    """
    log = structlog.get_logger(__name__)
    bs_model = read_bigstitcher_xml(os.path.join(alignment_url, "bigstitcher.xml"))
    interest_points_url = os.path.join(alignment_url, 'interestpoints.n5/')
    all_matches = fetch_all_matches(interest_points_url)
    valid_matches = {}
    for key, value in all_matches.items():
        if isinstance(value, BaseException):
            msg = f'An exception occurred when accessing {key}: {value.msg}'
            log.info(msg)
        else:
            valid_matches[key] = value
    
    summarized = summarize_matches(bs_model=bs_model, matches_dict=valid_matches)
    if output == 'csv':
        click.echo(summarized.write_csv())
    else:
        raise ValueError(f'Format {output} is not recognized. Allowed values: ("csv",)')