"""
speq CLI — cache management, synthesis, diffing, and test running.

Usage:
    speq synthesize <file.py>         Synthesize all @ai functions
    speq diff <file.py>               Show what would change on re-synthesis
    speq test <file.py>               Run generated tests for all @ai functions
    speq cache list <file.py>         List cached functions
    speq cache clear [<file.py>]      Clear cache (all or for one file)
"""

import difflib
import json
import subprocess
import sys
import tempfile
from pathlib import Path

import click

from speq.core import (
    _CACHE_DIR_NAME,
    _DEFAULT_MODEL,
    _cache_prefix,
    _get_cache_dir,
    _load_from_disk_cache,
    _spec_hash,
    synthesize_from_spec,
    ConfidenceScores,
)
from speq.scanner import scan_file


# ---------------------------------------------------------------------------
# Root group
# ---------------------------------------------------------------------------

@click.group()
@click.version_option(package_name="speq")
def cli():
    """speq — the speqification is the code."""
    pass


# ---------------------------------------------------------------------------
# speq synthesize
# ---------------------------------------------------------------------------

@cli.command()
@click.argument("filepath", type=click.Path(exists=True))
@click.option("--model", default=None, help="Model to use for synthesis.")
@click.option("--force", is_flag=True, help="Re-synthesize even if cached.")
def synthesize(filepath, model, force):
    """Synthesize all @ai functions in a file."""
    specs = scan_file(filepath)
    
    if not specs:
        click.echo(f"No @ai functions found in {filepath}")
        return
    
    click.echo(f"Found {len(specs)} @ai function(s) in {filepath}\n")
    
    for spec in specs:
        click.echo(f"  {spec.name}{spec.signature}")
        
        result = synthesize_from_spec(
            func_name=spec.name,
            signature=spec.signature,
            docstring=spec.docstring,
            source_file=spec.source_file,
            model=model,
            force=force,
        )
        
        if result.cached:
            click.echo(f"    ✓ cached (model: {result.model})")
        else:
            click.echo(f"    ✓ synthesized (model: {result.model})")
        
        if result.scores and result.scores.spec_clarity >= 0:
            sc = result.scores
            click.echo(f"    spec clarity: {sc.spec_clarity:.2f}  "
                       f"impl confidence: {sc.implementation_confidence:.2f}")
            if sc.assumptions:
                for a in sc.assumptions:
                    click.echo(f"    ⚠ assumption: {a}")
        
        if result.tests:
            click.echo(f"    tests: generated")
        
        click.echo()
    
    click.echo(f"Done. Cache: {_get_cache_dir(filepath)}")


# ---------------------------------------------------------------------------
# speq diff
# ---------------------------------------------------------------------------

@cli.command()
@click.argument("filepath", type=click.Path(exists=True))
@click.option("--model", default=None, help="Model to use for re-synthesis.")
def diff(filepath, model):
    """Show what would change if @ai functions were re-synthesized."""
    specs = scan_file(filepath)
    
    if not specs:
        click.echo(f"No @ai functions found in {filepath}")
        return
    
    cache_dir = _get_cache_dir(filepath)
    any_changes = False
    
    for spec in specs:
        new_sid = _spec_hash(spec.name, spec.signature, spec.docstring)
        
        # check if current spec hash matches what's cached
        cached_code = _load_from_disk_cache(cache_dir, spec.name, new_sid)
        
        if cached_code is not None:
            # spec hash unchanged — docstring hasn't been edited
            click.echo(f"  {spec.name}: spec unchanged, skipping")
            continue
        
        # spec hash changed or not cached — find the old cached version
        old_code = None
        old_files = sorted(cache_dir.glob(f"{spec.name}_*.py"))
        # exclude test files
        old_files = [f for f in old_files if not f.name.endswith("_test.py")]
        if old_files:
            old_code = old_files[0].read_text()
        
        if old_code is None:
            click.echo(f"  {spec.name}: not cached, would synthesize fresh")
            any_changes = True
            continue
        
        # re-synthesize with new spec
        any_changes = True
        click.echo(f"\n  {spec.name}: spec changed, re-synthesizing...")
        
        result = synthesize_from_spec(
            func_name=spec.name,
            signature=spec.signature,
            docstring=spec.docstring,
            source_file=spec.source_file,
            model=model,
            force=True,
        )
        
        old_lines = old_code.splitlines(keepends=True)
        new_lines = result.code.splitlines(keepends=True)
        
        diff_lines = list(difflib.unified_diff(
            old_lines, new_lines,
            fromfile=f"{spec.name} (cached)",
            tofile=f"{spec.name} (new)",
        ))
        
        if diff_lines:
            if result.scores and result.scores.spec_clarity >= 0:
                click.echo(f"    spec clarity: {result.scores.spec_clarity:.2f}  "
                           f"impl confidence: {result.scores.implementation_confidence:.2f}")
            
            for line in diff_lines:
                if line.startswith("+") and not line.startswith("+++"):
                    click.echo(click.style(f"    {line.rstrip()}", fg="green"))
                elif line.startswith("-") and not line.startswith("---"):
                    click.echo(click.style(f"    {line.rstrip()}", fg="red"))
                elif line.startswith("@@"):
                    click.echo(click.style(f"    {line.rstrip()}", fg="cyan"))
                else:
                    click.echo(f"    {line.rstrip()}")
            
            if click.confirm(f"\n  Accept new implementation for {spec.name}?"):
                # already saved by synthesize_from_spec with force=True
                # clean up old cache files
                for f in old_files:
                    f.unlink(missing_ok=True)
                    meta = f.with_suffix(".meta.json")
                    meta.unlink(missing_ok=True)
                    test = f.parent / f"{f.stem}_test.py"
                    test.unlink(missing_ok=True)
                click.echo(f"    ✓ updated")
            else:
                # remove the new cache, keep old
                new_prefix = _cache_prefix(spec.name, new_sid)
                (cache_dir / f"{new_prefix}.py").unlink(missing_ok=True)
                (cache_dir / f"{new_prefix}.meta.json").unlink(missing_ok=True)
                (cache_dir / f"{new_prefix}_test.py").unlink(missing_ok=True)
                click.echo(f"    ✗ kept old version")
        else:
            click.echo(f"  {spec.name}: spec changed but implementation is identical")
    
    if not any_changes:
        click.echo("\nAll implementations unchanged.")


# ---------------------------------------------------------------------------
# speq test
# ---------------------------------------------------------------------------

@cli.command()
@click.argument("filepath", type=click.Path(exists=True))
@click.option("--func", default=None, help="Test a specific function only.")
def test(filepath, func):
    """Run generated tests for @ai functions."""
    specs = scan_file(filepath)
    
    if not specs:
        click.echo(f"No @ai functions found in {filepath}")
        return
    
    if func:
        specs = [s for s in specs if s.name == func]
        if not specs:
            click.echo(f"Function '{func}' not found in {filepath}")
            return
    
    cache_dir = _get_cache_dir(filepath)
    total_passed = 0
    total_failed = 0
    total_skipped = 0
    
    for spec in specs:
        sid = _spec_hash(spec.name, spec.signature, spec.docstring)
        prefix = _cache_prefix(spec.name, sid)
        
        impl_file = cache_dir / f"{prefix}.py"
        test_file = cache_dir / f"{prefix}_test.py"
        
        if not impl_file.exists():
            click.echo(f"  {spec.name}: not synthesized yet, skipping")
            total_skipped += 1
            continue
        
        if not test_file.exists():
            click.echo(f"  {spec.name}: no tests generated, skipping")
            total_skipped += 1
            continue
        
        # create a temp file that patches the import
        impl_code = impl_file.read_text()
        test_code = test_file.read_text()
        
        # replace the placeholder import with the actual implementation
        test_code = test_code.replace(
            f"from __function__ import {spec.name}",
            impl_code  # inline the function definition
        )
        
        # also handle other common import patterns the LLM might generate
        test_code = test_code.replace(
            f"from {spec.name} import {spec.name}",
            impl_code
        )
        test_code = test_code.replace(
            f"import {spec.name}",
            impl_code
        )
        
        # write to temp file and run with pytest
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", prefix=f"speq_test_{spec.name}_",
            delete=False
        ) as f:
            f.write(test_code)
            temp_path = f.name
        
        try:
            result = subprocess.run(
                [sys.executable, "-m", "pytest", temp_path, "-v", "--tb=short", "--no-header"],
                capture_output=True, text=True, timeout=30,
            )
            
            # count results
            output = result.stdout
            passed = output.count(" PASSED")
            failed = output.count(" FAILED")
            
            if failed == 0 and passed > 0:
                click.echo(click.style(
                    f"  {spec.name}: {passed} tests passed ✓", fg="green"
                ))
            elif failed > 0:
                click.echo(click.style(
                    f"  {spec.name}: {passed} passed, {failed} failed ✗", fg="red"
                ))
                # show failures
                for line in output.split("\n"):
                    if "FAILED" in line or "AssertionError" in line or "assert" in line.lower():
                        click.echo(f"    {line.strip()}")
            else:
                click.echo(f"  {spec.name}: no test results")
            
            total_passed += passed
            total_failed += failed
            
        except subprocess.TimeoutExpired:
            click.echo(f"  {spec.name}: tests timed out")
            total_failed += 1
        except Exception as e:
            click.echo(f"  {spec.name}: error running tests: {e}")
            total_failed += 1
        finally:
            Path(temp_path).unlink(missing_ok=True)
    
    click.echo(f"\nTotal: {total_passed} passed, {total_failed} failed, {total_skipped} skipped")


# ---------------------------------------------------------------------------
# speq cache
# ---------------------------------------------------------------------------

@cli.group()
def cache():
    """Manage the synthesis cache."""
    pass


@cache.command("list")
@click.argument("filepath", type=click.Path(exists=True))
def cache_list(filepath):
    """List cached @ai functions for a file."""
    cache_dir = Path(filepath).parent / _CACHE_DIR_NAME
    
    if not cache_dir.exists():
        click.echo(f"No cache found for {filepath}")
        return
    
    meta_files = sorted(cache_dir.glob("*.meta.json"))
    
    if not meta_files:
        click.echo(f"Cache is empty: {cache_dir}")
        return
    
    click.echo(f"Cache: {cache_dir}\n")
    
    for meta_file in meta_files:
        try:
            data = json.loads(meta_file.read_text())
            func_name = data.get("function", "unknown")
            model = data.get("model", "unknown")
            spec_id = data.get("spec_id", "?")
            
            # check if implementation file exists
            impl_file = meta_file.with_suffix("").with_suffix(".py")
            impl_exists = impl_file.exists()
            impl_size = impl_file.stat().st_size if impl_exists else 0
            
            # check if test file exists
            test_file = meta_file.parent / f"{meta_file.stem.replace('.meta', '')}_test.py"
            has_tests = test_file.exists()
            
            scores = data.get("scores", {})
            sc = scores.get("spec_clarity", -1)
            ic = scores.get("implementation_confidence", -1)
            
            score_str = ""
            if sc >= 0:
                score_str = f"  clarity: {sc:.2f}  confidence: {ic:.2f}"
            
            tests_str = "  tests: ✓" if has_tests else ""
            
            click.echo(
                f"  {func_name:30s}  model: {model or 'unknown':30s}"
                f"{score_str}{tests_str}"
            )
        except (json.JSONDecodeError, KeyError):
            click.echo(f"  {meta_file.name}: corrupted metadata")


@cache.command("clear")
@click.argument("filepath", type=click.Path(exists=True), required=False)
@click.option("--func", default=None, help="Clear cache for a specific function only.")
def cache_clear(filepath, func):
    """Clear the synthesis cache."""
    if filepath:
        cache_dir = Path(filepath).parent / _CACHE_DIR_NAME
    else:
        cache_dir = Path(".") / _CACHE_DIR_NAME
    
    if not cache_dir.exists():
        click.echo("No cache to clear.")
        return
    
    if func:
        # clear only files matching this function name
        removed = 0
        for f in cache_dir.glob(f"{func}_*"):
            f.unlink()
            removed += 1
        click.echo(f"Removed {removed} cache file(s) for '{func}'")
    else:
        # clear everything
        import shutil
        file_count = sum(1 for _ in cache_dir.iterdir())
        shutil.rmtree(cache_dir)
        click.echo(f"Cleared {file_count} file(s) from {cache_dir}")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    cli()


if __name__ == "__main__":
    main()
