from __future__ import annotations
import json
import click
from .table_encoder import encode_csv, EncodeConfig
from .table_decoder import decode_to_csv
from .metrics import eval_table
import pandas as pd

# Image compression imports
from .image_encoder import encode_image, ImageEncodeConfig
from .image_decoder import decode_to_image_file, decode_image
from .image_metrics import eval_image
import numpy as np
from PIL import Image


@click.group()
def main():
    """Sempress: semantic compression for tables and images"""
    pass


# ========== TABLE COMMANDS ==========

@main.command()
@click.option("--in", "in_path", required=True, type=click.Path(exists=True))
@click.option("--out", "out_path", required=True, type=click.Path())
@click.option("--lock-cols", default="", help="Comma-separated columns to store losslessly")
@click.option("--residual-cols", default="", help="Comma-separated columns to store residuals for")
@click.option("--k", default=64, show_default=True, type=int, help="Codebook size per numeric column")
@click.option("--uncert-thresh", default=0.2, show_default=True, type=float, help="Relative error threshold for uncertainty mask")
@click.option("--seed", default=42, show_default=True, type=int)
def encode(in_path, out_path, lock_cols, residual_cols, k, uncert_thresh, seed):
    """Encode a CSV file to .smp format"""
    cfg = EncodeConfig(
        lock_cols=[c for c in lock_cols.split(",") if c.strip()],
        residual_cols=[c for c in residual_cols.split(",") if c.strip()],
        k=k,
        uncertainty_thresh=uncert_thresh,
        random_state=seed,
    )
    blob = encode_csv(in_path, cfg)
    with open(out_path, "wb") as f:
        f.write(blob)
    click.echo(f"Wrote {out_path}")

@main.command()
@click.option("--in", "in_path", required=True, type=click.Path(exists=True))
@click.option("--out", "out_path", required=True, type=click.Path())
def decode(in_path, out_path):
    """Decode a .smp file back to CSV"""
    meta = decode_to_csv(in_path, out_path)
    click.echo(f"Reconstructed → {out_path}")
    click.echo("Uncertainty summary:" )
    click.echo(json.dumps(meta.get("uncertain", {}), indent=2))

@main.command()
@click.option("--original", required=True, type=click.Path(exists=True))
@click.option("--recon", required=True, type=click.Path(exists=True))
@click.option("--lock-cols", default="", help="Comma-separated locked columns")
def eval(original, recon, lock_cols):
    """Evaluate reconstruction quality for tables"""
    df = pd.read_csv(original)
    dfh = pd.read_csv(recon)
    locks = [c for c in lock_cols.split(",") if c.strip()]
    metrics = eval_table(df, dfh, locks)
    click.echo(json.dumps(metrics, indent=2))


# ========== IMAGE COMMANDS ==========

@main.command(name="encode-image")
@click.option("--in", "in_path", required=True, type=click.Path(exists=True), help="Input image file")
@click.option("--out", "out_path", required=True, type=click.Path(), help="Output .smp file")
@click.option("--patch-size", default=8, show_default=True, type=click.Choice(["8", "16"]), help="Patch size (8x8 or 16x16)")
@click.option("--k", default=256, show_default=True, type=int, help="Codebook size per channel")
@click.option("--color-space", default="YCbCr", show_default=True, type=click.Choice(["RGB", "YCbCr", "Grayscale"]), help="Color space for compression")
@click.option("--residual-thresh", default=0.1, show_default=True, type=float, help="Error threshold for residuals (0-1)")
@click.option("--quality", default="balanced", show_default=True, type=click.Choice(["fast", "balanced", "high"]), help="Quality mode")
@click.option("--seed", default=42, show_default=True, type=int)
def encode_image_cmd(in_path, out_path, patch_size, k, color_space, residual_thresh, quality, seed):
    """Encode an image file to .smp format"""
    cfg = ImageEncodeConfig(
        patch_size=int(patch_size),
        k=k,
        color_space=color_space,
        residual_threshold=residual_thresh,
        quality_mode=quality,
        random_state=seed,
    )
    
    click.echo(f"Compressing {in_path}...")
    click.echo(f"  Patch size: {patch_size}×{patch_size}")
    click.echo(f"  Codebook size: {k}")
    click.echo(f"  Color space: {color_space}")
    
    blob = encode_image(in_path, cfg)
    
    with open(out_path, "wb") as f:
        f.write(blob)
    
    # Calculate stats
    original_size = Image.open(in_path).size
    img = np.array(Image.open(in_path))
    ratio = img.nbytes / len(blob)
    
    click.echo(f"\n✓ Compressed: {img.nbytes:,} bytes → {len(blob):,} bytes")
    click.echo(f"  Compression ratio: {ratio:.2f}×")
    click.echo(f"  Saved to: {out_path}")

@main.command(name="decode-image")
@click.option("--in", "in_path", required=True, type=click.Path(exists=True), help="Input .smp file")
@click.option("--out", "out_path", required=True, type=click.Path(), help="Output image file")
def decode_image_cmd(in_path, out_path):
    """Decode a .smp file back to an image"""
    click.echo(f"Decompressing {in_path}...")
    
    decode_to_image_file(open(in_path, "rb").read(), out_path)
    
    click.echo(f"✓ Reconstructed → {out_path}")

@main.command(name="eval-image")
@click.option("--original", required=True, type=click.Path(exists=True), help="Original image file")
@click.option("--recon", required=True, type=click.Path(exists=True), help="Reconstructed image file")
@click.option("--compressed", required=False, type=click.Path(exists=True), help="Compressed .smp file (optional)")
def eval_image_cmd(original, recon, compressed):
    """Evaluate image reconstruction quality"""
    click.echo("Evaluating image quality...")
    
    # Load images
    orig_img = np.array(Image.open(original))
    recon_img = np.array(Image.open(recon))
    
    # Get compressed size
    if compressed:
        compressed_size = open(compressed, "rb").seek(0, 2)
    else:
        compressed_size = orig_img.nbytes  # Fallback
    
    # Calculate metrics
    metrics = eval_image(orig_img, recon_img, compressed_size)
    
    click.echo("\n" + "="*50)
    click.echo("QUALITY METRICS")
    click.echo("="*50)
    click.echo(f"PSNR:              {metrics['psnr_db']:.2f} dB")
    click.echo(f"SSIM:              {metrics['ssim']:.4f}")
    click.echo(f"MSE:               {metrics['mse']:.2f}")
    click.echo(f"MAE:               {metrics['mae']:.2f}")
    click.echo(f"Compression ratio: {metrics['compression_ratio']:.2f}×")
    click.echo(f"Bits per pixel:    {metrics['bits_per_pixel']:.3f}")
    click.echo(f"Size reduction:    {metrics['size_reduction_percent']:.1f}%")
    
    from .image_metrics import quality_grade
    grade = quality_grade(metrics['psnr_db'], metrics['ssim'])
    click.echo(f"\nQuality grade: {grade}")
    click.echo("="*50)
