"""bandito setup — configure trace storage."""

from __future__ import annotations

import os


def run_setup_config() -> None:
    """Interactive setup: where are your traces?"""
    from ..config import load_config, save_config

    config = load_config()

    print(f"\n  bandito setup config\n")
    print("  Where are your traces?\n")
    print("  1. Local (default — ~/.bandito/traces)")
    print("  2. S3")
    print("  3. PostgreSQL")
    print("  4. Langfuse (read from existing Langfuse project)")
    print()

    try:
        choice = input("  [1]: ").strip() or "1"
    except (EOFError, KeyboardInterrupt):
        print("\n  Cancelled.\n")
        return

    if choice == "1":
        _setup_local(config)
    elif choice == "2":
        _setup_s3(config)
    elif choice == "3":
        _setup_postgres(config)
    elif choice == "4":
        _setup_langfuse(config)
    else:
        print(f"  Unknown option: {choice}\n")
        return


def _setup_local(config) -> None:
    from ..config import save_config

    try:
        directory = input(f"  Traces directory [{config.traces_dir}]: ").strip()
    except (EOFError, KeyboardInterrupt):
        print("\n  Cancelled.\n")
        return

    if directory:
        config.traces_dir = directory
    config.storage.backend = "local"

    save_config(config)
    print(f"\n  Storage set to local.")
    print(f"  Traces directory: {config.traces_dir}\n")


def _setup_s3(config) -> None:
    from ..config import save_config

    try:
        bucket = input(f"  S3 bucket [{config.storage.s3.bucket}]: ").strip()
        if not bucket:
            bucket = config.storage.s3.bucket
        prefix = input(f"  S3 prefix [{config.storage.s3.prefix}]: ").strip()
        if not prefix:
            prefix = config.storage.s3.prefix
        region = input(f"  Region [{config.storage.s3.region}]: ").strip() or config.storage.s3.region
        endpoint = input(f"  Endpoint (enter for AWS default) [{config.storage.s3.endpoint}]: ").strip()
        if not endpoint:
            endpoint = config.storage.s3.endpoint
    except (EOFError, KeyboardInterrupt):
        print("\n  Cancelled.\n")
        return

    if not bucket:
        print("  Bucket is required.\n")
        return
    try:
        import boto3
        s3 = boto3.client(
            "s3",
            region_name=region,
            endpoint_url=endpoint or None,
            aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
            aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
        )
        s3.head_bucket(Bucket=bucket)
        print("  Connected.\n")
    except ImportError:
        print("  Warning: boto3 not installed. Skipping validation.\n")
    except Exception as e:
        print(f"  Could not connect: {e}")
        print("  Check your AWS credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY).\n")
        return

    config.storage.backend = "s3"
    config.storage.s3.bucket = bucket
    config.storage.s3.prefix = prefix
    config.storage.s3.region = region
    if endpoint:
        config.storage.s3.endpoint = endpoint

    save_config(config)
    print(f"  Storage set to S3.")
    print(f"    Bucket: {bucket}")
    if prefix:
        print(f"    Prefix: {prefix}")
    print()

    # Also save to backend if we have an API key
    _save_org_storage("s3", {
        "bucket": bucket, "region": region,
        **({"prefix": prefix} if prefix else {}),
        **({"endpoint": endpoint} if endpoint else {}),
    })


def _setup_postgres(config) -> None:
    from ..config import save_config

    try:
        url = input(f"  PostgreSQL URL [{config.storage.postgres.url}]: ").strip()
        if not url:
            url = config.storage.postgres.url
    except (EOFError, KeyboardInterrupt):
        print("\n  Cancelled.\n")
        return

    if not url:
        print("  URL is required.\n")
        return
    print("\n  Validating PostgreSQL connectivity...")
    try:
        import psycopg
        conn = psycopg.connect(url, autocommit=True)
        conn.close()
        print("  Connected.\n")
    except ImportError:
        print("  Warning: psycopg not installed. Skipping validation.\n")
    except Exception as e:
        print(f"  Could not connect: {e}\n")
        return

    config.storage.backend = "postgres"
    config.storage.postgres.url = url

    save_config(config)
    print(f"  Storage set to PostgreSQL.\n")

    _save_org_storage("postgres", {"url": url})


def _setup_langfuse(config) -> None:
    from ..config import save_config

    print()
    print("  Langfuse reads traces from your existing Langfuse project.")
    print("  Bandito adds grading, judging, and analysis on top.\n")

    try:
        public_key_input = input(f"  Langfuse Public Key (pk-lf-...) [{config.storage.langfuse.public_key}]: ").strip()
        public_key = public_key_input or config.storage.langfuse.public_key
        secret_key_input = input(f"  Langfuse Secret Key (sk-lf-...) [{config.storage.langfuse.secret_key[:8]}...]: ").strip()
        secret_key = secret_key_input or config.storage.langfuse.secret_key
        base_url_input = input(f"  Langfuse Base URL [{config.storage.langfuse.base_url}]: ").strip()
        base_url = base_url_input or config.storage.langfuse.base_url
    except (EOFError, KeyboardInterrupt):
        print("\n  Cancelled.\n")
        return

    if not public_key:
        print("  Public key is required.\n")
        return
    if not secret_key:
        print("  Secret key is required.\n")
        return

    # Validate
    print("\n  Validating Langfuse connectivity...")
    try:
        from ..storage.langfuse import LangfuseStore
        store = LangfuseStore(
            public_key=public_key,
            secret_key=secret_key,
            base_url=base_url,
        )
        err = store.validate()
        if err:
            print(f"  {err}\n")
            return
        print("  Connected.\n")

        projects = store.list_projects()
        if projects:
            print(f"  Found {len(projects)} trace names:")
            for p in projects[:10]:
                print(f"    {p}")
            if len(projects) > 10:
                print(f"    ... and {len(projects) - 10} more")
            print()
        else:
            print("  No traces found yet. Send some traces via Langfuse first.\n")
    except ImportError:
        print("  langfuse not installed. Run: pip install bandito[langfuse]\n")
        return
    except Exception as e:
        print(f"  Could not connect: {e}\n")
        return

    # Save to config.toml
    config.storage.backend = "langfuse"
    config.storage.langfuse.public_key = public_key
    config.storage.langfuse.secret_key = secret_key
    config.storage.langfuse.base_url = base_url

    save_config(config)
    print(f"  Saved to ~/.bandito/config.toml\n")

    print("  You're set. Try:")
    if projects:
        example = projects[0]
        print(f"    bandito observe projects")
        print(f"    bandito diagnostics --project {example}")
        print(f"    bandito analyze tradeoffs --project {example}")
    else:
        print(f"    bandito observe projects")
    print()


def _save_org_storage(backend: str, storage_config: dict) -> None:
    """Try to save org storage to backend. Silent if no API key."""
    from . import get_client

    client = get_client()
    if client is None:
        return
    try:
        client.save_org_storage(backend, storage_config)
    except Exception:
        pass
