from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import typer

from .logging_utils import configure_logging
from .data_registry import DataRegistry
from .download import download_files, DownloadTarget
from .io import export_rda_to_xlsx, read_xlsx_tables, build_xlsx_from_repo
from .country_workbook import build_country_only_workbook
from .data_readers import from_xlsx_tables
from .demo import DemoParams, run_demo_and_export, compare_intermediates
from .psa import LHSSpec, ParameterDist, run_psa
import numpy as np
from scipy.stats import beta as sp_beta


app = typer.Typer(add_completion=False, no_args_is_help=True)


@app.command()
def init_registry(base_dir: str = typer.Argument(..., help="Base repository directory containing data/")) -> None:
    logger = configure_logging()
    reg = DataRegistry(Path(base_dir))
    reg.default_from_repo()
    # Print confirmation
    logger.info("Initialized data registry from repository data directory")


@app.command()
def register_xlsx(base_dir: str = typer.Argument(...), year: str = typer.Argument(...), xlsx: str = typer.Argument(...)) -> None:
    """Register an xlsx bundle for a given dataset year (2012/2018/2020)."""
    from .data_registry import DatasetPaths

    reg = DataRegistry(Path(base_dir))
    reg.register(year, DatasetPaths(rda_bundle=None, xlsx_bundle=Path(xlsx)))  # type: ignore[arg-type]
    typer.echo(f"Registered {year} -> {xlsx}")


@app.command()
def show_versions(base_dir: str = typer.Argument(...)) -> None:
    reg = DataRegistry(Path(base_dir))
    versions = reg.list_versions()
    typer.echo(json.dumps(versions, indent=2))


@app.command()
def data_download(config: str = typer.Argument(..., help="JSON mapping names to {url, out}")) -> None:
    """Download current data files using a JSON config.

    Example config:
    {
      "globocan2020_incidence": {"url": "https://.../incidence2020.csv", "out": "data/raw/incidence2020.csv"}
    }
    """
    logger = configure_logging()
    cfg_path = Path(config)
    spec = json.loads(cfg_path.read_text())
    targets = {
        name: DownloadTarget(url=item["url"], out=Path(item["out"])) for name, item in spec.items()
    }
    download_files(targets)
    logger.info("Downloads completed")


@app.command()
def data_convert_rda_to_xlsx(rda_path: str = typer.Argument(...), xlsx_path: str = typer.Argument(...)) -> None:
    """Convert an R .rda bundle to an editable .xlsx workbook with multiple sheets."""
    export_rda_to_xlsx(Path(rda_path), Path(xlsx_path))
    typer.echo(f"Wrote {xlsx_path}")


@app.command()
def data_validate_xlsx(xlsx_path: str = typer.Argument(...), year: str = typer.Argument(...)) -> None:
    """Validate required sheets exist and show basic shapes for a given dataset year."""
    tables = read_xlsx_tables(Path(xlsx_path))
    data = from_xlsx_tables(tables, year)  # type: ignore[arg-type]
    summary = {
        "incidence": list(data.incidence.columns),
        "mortality": list(data.mortality.columns),
        "prevalence5y": list(data.prevalence5y.columns),
        "sexual_debut": list(data.sexual_debut.columns),
        "hpv_distribution": list(data.hpv_distribution.columns),
        "disability_weights": list(data.disability_weights.columns),
        "population": list(data.population.columns),
        "unwpp_nqx": list(data.unwpp_nqx.columns),
    }
    typer.echo(json.dumps({k: len(v) for k, v in summary.items()}, indent=2))


@app.command()
def data_build_xlsx_from_repo(base_dir: str = typer.Argument(...), year: str = typer.Argument(...), out: str = typer.Argument(...)) -> None:
    """Build a consolidated .xlsx from the repo's .rda data files for a dataset year."""
    build_xlsx_from_repo(Path(base_dir) / "data", year, Path(out))
    typer.echo(f"Wrote {out}")


@app.command()
def data_country_xlsx(
    base_dir: str = typer.Argument(...),
    year: str = typer.Argument(...),
    iso3: str = typer.Argument(...),
) -> None:
    """Create a country-only workbook (e.g., TUR) for the given dataset year.

    Writes data/globocan_{year}_{iso3}.xlsx and refreshes the registry so it is preferred.
    """
    path = build_country_only_workbook(Path(base_dir), year, iso3)
    typer.echo(f"Wrote {path}")


@app.command()
def demo_run(
    base_dir: str = typer.Argument(...),
    country: str = typer.Argument(...),
    dataset_year: str = typer.Argument(...),
    age_vaccination: int = typer.Argument(...),
    current_year: int = typer.Argument(...),
    oldest_year: int = typer.Argument(...),
    next_n_years: int = typer.Argument(...),
    coverage: float = typer.Argument(...),
    out_dir: str = typer.Argument(...),
    vaccine: str = typer.Option("4vHPV"),
) -> None:
    params = DemoParams(
        country_iso3=country,
        dataset_year=dataset_year,
        vaccine=vaccine,  # type: ignore[arg-type]
        age_vaccination=age_vaccination,
        current_year=current_year,
        oldest_year=oldest_year,
        next_n_years=next_n_years,
        coverage_vac=coverage,
    )
    run_demo_and_export(Path(base_dir), params, Path(out_dir))
    typer.echo("Demo intermediates exported (python)")


@app.command()
def demo_compare(r_dir: str = typer.Argument(...), p_dir: str = typer.Argument(...)) -> None:
    diffs = compare_intermediates(Path(r_dir), Path(p_dir))
    typer.echo(json.dumps(diffs, indent=2))


@app.command()
def psa_run(
    size: int = typer.Argument(...),
    seed: int = typer.Option(1),
) -> None:
    # Example PSA over two parameters with beta distributions
    dists = {
        "incidence_ratio": ParameterDist(inv_cdf=lambda u: sp_beta(a=2, b=5).ppf(u)),
        "mortality_ratio": ParameterDist(inv_cdf=lambda u: sp_beta(a=2, b=5).ppf(u)),
    }

    def simulate(params: dict) -> dict:
        # placeholder output; wire into country simulation later
        return {"outcome": float(params["incidence_ratio"] * 1000.0)}

    results = run_psa(LHSSpec(size=size, dims=len(dists), seed=seed), dists, simulate)
    typer.echo(json.dumps(results[:5], indent=2))