Analyzing Data

uv.lock

"""Persistent cache for concepts, patterns, and table schemas.

Cache files are stored at ~/.astro/ai/cache/:
- concepts.json: concept → table mapping (e.g., "customers" → "HQ.MODEL.ORGS")
- patterns.json: question type → query strategy
- tables.json: table schema cache (columns, types, row counts)
"""

import json
from datetime import datetime, timedelta
from pathlib import Path

CACHE_DIR = Path.home() / ".astro" / "ai" / "cache"

# Default TTL for cache entries
DEFAULT_TTL_DAYS = 90


def _ensure_cache_dir():
    CACHE_DIR.mkdir(parents=True, exist_ok=True)


def _load_json(filename: str) -> dict:
    path = CACHE_DIR / filename
    if path.exists():
        return json.loads(path.read_text())
    return {}


def _save_json(filename: str, data: dict):
    _ensure_cache_dir()
    path = CACHE_DIR / filename
    path.write_text(json.dumps(data, indent=2, default=str))


# --- Concept Cache ---


def lookup_concept(concept: str) -> dict | None:
    """Look up a concept (e.g., 'customers') to find its table."""
    concepts = _load_json("concepts.json")
    return concepts.get(concept.lower().strip())


def learn_concept(
    concept: str,
    table: str,
    key_column: str | None = None,
    date_column: str | None = None,
):
    """Store a concept -> table mapping for future use."""
    concepts = _load_json("concepts.json")
    concepts[concept.lower().strip()] = {
        "table": table,
        "key_column": key_column,
        "date_column": date_column,
        "learned_at": datetime.now().isoformat(),
    }
    _save_json("concepts.json", concepts)
    return concepts[concept.lower().strip()]


def list_concepts() -> dict:
    """List all learned concepts."""
    return _load_json("concepts.json")


# --- Pattern Cache ---


def lookup_pattern(question: str) -> list[dict]:
    """Find patterns that match a question. Returns list of matching patterns."""
    patterns = _load_json("patterns.json")
    question_lower = question.lower()
    matches = []
    for name, pattern in patterns.items():
        for qtype in pattern.get("question_types", []):
            keywords = qtype.lower().replace("x", "").split()
            if all(kw in question_lower for kw in keywords if len(kw) > 2):
                matches.append({"name": name, **pattern})
                break
    return sorted(matches, key=lambda p: p.get("success_count", 0), reverse=True)


def learn_pattern(
    name: str,
    question_types: list[str],
    strategy: list[str],
    tables_used: list[str],
    gotchas: list[str],
    example_query: str | None = None,
):
    """Store a query pattern/strategy for a type of question."""
    patterns = _load_json("patterns.json")
    patterns[name.lower().strip()] = {
        "question_types": question_types,
        "strategy": strategy,
        "tables_used": tables_used,
        "gotchas": gotchas,
        "example_query": example_query,
        "created_at": datetime.now().isoformat(),
        "success_count": 1,
        "failure_count": 0,
    }
    _save_json("patterns.json", patterns)
    return patterns[name.lower().strip()]


def record_pattern_outcome(name: str, success: bool):
    """Record whether a pattern helped or failed."""
    patterns = _load_json("patterns.json")
    key = name.lower().strip()
    if key in patterns:
        if success:
            patterns[key]["success_count"] = patterns[key].get("success_count", 0) + 1
        else:
            patterns[key]["failure_count"] = patterns[key].get("failure_count", 0) + 1
        _save_json("patterns.json", patterns)
        return patterns[key]
    return None


def list_patterns() -> dict:
    """List all learned patterns."""
    return _load_json("patterns.json")


def delete_pattern(name: str) -> bool:
    """Delete a pattern by name. Returns True if it existed."""
    patterns = _load_json("patterns.json")
    key = name.lower().strip()
    if key in patterns:
        del patterns[key]
        _save_json("patterns.json", patterns)
        return True
    return False


# --- Cache Management ---


def _is_stale(learned_at: str, ttl_days: int = DEFAULT_TTL_DAYS) -> bool:
    """Check if an entry is older than TTL."""
    try:
        learned = datetime.fromisoformat(learned_at)
        return datetime.now() - learned > timedelta(days=ttl_days)
    except (ValueError, TypeError):
        return False


def cache_stats() -> dict:
    """Get cache statistics."""
    concepts = _load_json("concepts.json")
    patterns = _load_json("patterns.json")

    stale_concepts = sum(
        1 for c in concepts.values() if _is_stale(c.get("learned_at", ""))
    )
    stale_patterns = sum(
        1 for p in patterns.values() if _is_stale(p.get("created_at", ""))
    )

    return {
        "concepts_count": len(concepts),
        "patterns_count": len(patterns),
        "stale_concepts": stale_concepts,
        "stale_patterns": stale_patterns,
        "cache_dir": str(CACHE_DIR),
        "ttl_days": DEFAULT_TTL_DAYS,
    }


def clear_cache(cache_type: str = "all", purge_stale_only: bool = False) -> dict:
    """Clear cache entries.

    Args:
        cache_type: "all", "concepts", or "patterns"
        purge_stale_only: If True, only remove entries older than TTL

    Returns:
        Summary of what was cleared
    """
    result = {"concepts_cleared": 0, "patterns_cleared": 0}

    if cache_type in ("all", "concepts"):
        concepts = _load_json("concepts.json")
        if purge_stale_only:
            original = len(concepts)
            concepts = {
                k: v
                for k, v in concepts.items()
                if not _is_stale(v.get("learned_at", ""))
            }
            result["concepts_cleared"] = original - len(concepts)
            _save_json("concepts.json", concepts)
        else:
            result["concepts_cleared"] = len(concepts)
            _save_json("concepts.json", {})

    if cache_type in ("all", "patterns"):
        patterns = _load_json("patterns.json")
        if purge_stale_only:
            original = len(patterns)
            patterns = {
                k: v
                for k, v in patterns.items()
                if not _is_stale(v.get("created_at", ""))
            }
            result["patterns_cleared"] = original - len(patterns)
            _save_json("patterns.json", patterns)
        else:
            result["patterns_cleared"] = len(patterns)
            _save_json("patterns.json", {})

    return result


# --- Table Schema Cache ---


def get_table(full_name: str) -> dict | None:
    """Get cached table schema by full name (DATABASE.SCHEMA.TABLE)."""
    tables = _load_json("tables.json")
    return tables.get(full_name.upper())


def set_table(
    full_name: str,
    columns: list[dict],
    row_count: int | None = None,
    comment: str | None = None,
) -> dict:
    """Cache a table's schema.

    Args:
        full_name: Full table name (DATABASE.SCHEMA.TABLE)
        columns: List of column dicts [{name, type, nullable, comment}, ...]
        row_count: Optional row count
        comment: Optional table description

    Returns:
        The cached table entry
    """
    tables = _load_json("tables.json")
    entry = {
        "full_name": full_name.upper(),
        "columns": columns,
        "row_count": row_count,
        "comment": comment,
        "cached_at": datetime.now().isoformat(),
    }
    tables[full_name.upper()] = entry
    _save_json("tables.json", tables)
    return entry


def list_tables() -> dict:
    """List all cached table schemas."""
    return _load_json("tables.json")


def delete_table(full_name: str) -> bool:
    """Remove a table from cache. Returns True if it existed."""
    tables = _load_json("tables.json")
    key = full_name.upper()
    if key in tables:
        del tables[key]
        _save_json("tables.json", tables)
        return True
    return False


# --- Bulk Import ---


def load_concepts_from_warehouse_md(path: Path | None = None) -> int:
    """Parse warehouse.md and populate cache with Quick Reference entries.

    Looks for a markdown table with columns: Concept | Table | Key Column | Date Column

    Args:
        path: Path to warehouse.md. If None, searches common locations.

    Returns:
        Number of concepts loaded into cache.
    """
    import re

    # Find warehouse.md if not provided
    if path is None:
        locations = [
            Path(".astro/warehouse.md"),
            Path.home() / ".astro" / "agents" / "warehouse.md",
            Path("warehouse.md"),
        ]
        for loc in locations:
            if loc.exists():
                path = loc
                break

    if path is None or not path.exists():
        return 0

    content = path.read_text(encoding="utf-8")
    concepts_loaded = 0

    # Find markdown table rows: | concept | table | key_col | date_col |
    # Skip header rows (contain "Concept" or "---")
    table_pattern = re.compile(
        r"^\|\s*([^|]+)\s*\|\s*([^|]+)\s*\|(?:\s*([^|]*)\s*\|)?(?:\s*([^|]*)\s*\|)?",
        re.MULTILINE,
    )

    for match in table_pattern.finditer(content):
        concept = match.group(1).strip()
        table = match.group(2).strip()
        key_col = match.group(3).strip() if match.group(3) else None
        date_col = match.group(4).strip() if match.group(4) else None

        # Skip header/separator rows
        if not concept or concept.lower() == "concept" or "---" in concept:
            continue
        if not table or table.lower() == "table" or "---" in table:
            continue
        # Skip if table doesn't look valid (should have dots for fully qualified name)
        if "." not in table:
            continue

        # Normalize empty values
        if key_col in ("-", "", None):
            key_col = None
        if date_col in ("-", "", None):
            date_col = None

        learn_concept(concept, table, key_col, date_col)
        concepts_loaded += 1

    return concepts_loaded

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "click>=8.0.0",
#     "jupyter-client>=8.0.0",
#     "ipykernel>=6.0.0",
#     "pyyaml>=6.0",
#     "python-dotenv>=1.0.0",
#     "cryptography>=41.0.0",
# ]
# ///
"""CLI for the analyzing-data skill.

Usage:
    uv run scripts/cli.py start    # Start kernel with Snowflake
    uv run scripts/cli.py exec "df = run_sql('SELECT ...')"
    uv run scripts/cli.py status   # Check kernel status
    uv run scripts/cli.py stop     # Stop kernel
"""

import json
import shutil
import sys

import click

# Add parent directory to path for lib imports
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from kernel import KernelManager
from warehouse import WarehouseConfig
import cache


def check_uv_installed():
    """Check if uv is installed and provide helpful error if not."""
    if not shutil.which("uv"):
        click.echo("Error: uv is not installed.", err=True)
        click.echo(
            "Install with: curl -LsSf https://astral.sh/uv/install.sh | sh", err=True
        )
        sys.exit(1)


@click.group()
@click.version_option(version="0.1.0")
def main():
    """Jupyter kernel CLI for data analysis with Snowflake."""
    pass


@main.group()
def warehouse():
    """Manage warehouse connections."""


@warehouse.command("list")
def warehouse_list():
    """List available warehouse connections."""
    try:
        config = WarehouseConfig.load()
        if not config.connectors:
            click.echo("No warehouses configured")
            return

        default_name, _ = config.get_default()
        for name, conn in config.connectors.items():
            marker = " (default)" if name == default_name else ""
            click.echo(f"{name}: {conn.connector_type()}{marker}")
    except FileNotFoundError:
        click.echo("No warehouse config found at ~/.astro/agents/warehouse.yml")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)


@main.command()
@click.option("--warehouse", "-w", help="Warehouse name from config")
def start(warehouse: str | None):
    """Start kernel with Snowflake connection."""
    check_uv_installed()
    km = KernelManager()

    if km.is_running:
        click.echo("Kernel already running")
        return

    try:
        config = WarehouseConfig.load()
        wh_name, wh_config = (
            (warehouse, config.connectors[warehouse])
            if warehouse
            else config.get_default()
        )
        click.echo(f"Using warehouse: {wh_name}")
    except FileNotFoundError as e:
        click.echo(f"Error: {e}", err=True)
        click.echo(
            "Create ~/.astro/agents/warehouse.yml with your Snowflake credentials",
            err=True,
        )
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

    env_vars = wh_config.get_env_vars_for_kernel()
    extra_packages = wh_config.get_required_packages()

    km.start(env_vars=env_vars, extra_packages=extra_packages)

    result = km.execute(wh_config.to_python_prelude(), timeout=60.0)
    if not result.success:
        click.echo(f"Connection error:\n{result.error}", err=True)
        km.stop()
        sys.exit(1)
    click.echo(result.output)


@main.command("exec")
@click.argument("code")
@click.option("--timeout", "-t", default=30.0, help="Timeout in seconds")
def execute(code: str, timeout: float):
    """Execute Python code in the kernel. Auto-starts kernel if not running."""
    km = KernelManager()

    if not km.is_running:
        check_uv_installed()
        try:
            config = WarehouseConfig.load()
            wh_name, wh_config = config.get_default()
            click.echo(f"Starting kernel with: {wh_name}", err=True)
            env_vars = wh_config.get_env_vars_for_kernel()
            extra_packages = wh_config.get_required_packages()
            km.start(env_vars=env_vars, extra_packages=extra_packages)
            result = km.execute(wh_config.to_python_prelude(), timeout=60.0)
            if result.output:
                click.echo(result.output, err=True)
            if not result.success:
                click.echo(f"Connection error:\n{result.error}", err=True)
                km.stop()
                sys.exit(1)
        except Exception as e:
            click.echo(f"Error starting kernel: {e}", err=True)
            sys.exit(1)

    result = km.execute(code, timeout=timeout)
    if result.output:
        click.echo(result.output, nl=False)
    if result.error:
        click.echo(result.error, err=True)
        sys.exit(1)


@main.command()
def stop():
    """Stop the kernel."""
    KernelManager().stop()


@main.command()
def restart():
    """Restart the kernel (stop + start)."""
    km = KernelManager()
    km.stop()

    try:
        config = WarehouseConfig.load()
        wh_name, wh_config = config.get_default()
        click.echo(f"Restarting kernel with: {wh_name}")
        env_vars = wh_config.get_env_vars_for_kernel()
        extra_packages = wh_config.get_required_packages()
        km.start(env_vars=env_vars, extra_packages=extra_packages)
        result = km.execute(wh_config.to_python_prelude(), timeout=60.0)
        if result.output:
            click.echo(result.output)
        if not result.success:
            click.echo(f"Connection error:\n{result.error}", err=True)
            sys.exit(1)
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)


@main.command()
@click.option("--json", "as_json", is_flag=True)
def status(as_json: bool):
    """Check kernel status."""
    info = KernelManager().status()
    if as_json:
        click.echo(json.dumps(info, indent=2))
    else:
        if info["running"]:
            click.echo(
                f"Kernel: {'running' if info['responsive'] else 'running (unresponsive)'}"
            )
        else:
            click.echo("Kernel: not running")


@main.command("install")
@click.argument("packages", nargs=-1, required=True)
def install_packages(packages: tuple):
    """Install additional packages into the kernel environment.

    Example: uv run scripts/cli.py install plotly scipy
    """
    km = KernelManager()
    success, message = km.install_packages(list(packages))
    if success:
        click.echo(message)
    else:
        click.echo(f"Error: {message}", err=True)
        sys.exit(1)


@main.command()
def ensure():
    """Ensure kernel is running (start if needed). Used by hooks."""
    check_uv_installed()
    km = KernelManager()
    if km.is_running:
        return

    try:
        config = WarehouseConfig.load()
        wh_name, wh_config = config.get_default()
        click.echo(f"Starting kernel with: {wh_name}", err=True)
        env_vars = wh_config.get_env_vars_for_kernel()
        extra_packages = wh_config.get_required_packages()
        km.start(env_vars=env_vars, extra_packages=extra_packages)
        result = km.execute(wh_config.to_python_prelude(), timeout=60.0)
        if result.output:
            click.echo(result.output, err=True)
    except Exception as e:
        click.echo(f"Warning: {e}", err=True)
        km.start()


@main.group()
def concept():
    """Manage concept cache (concept -> table mappings)."""
    pass


@concept.command("lookup")
@click.argument("name")
def concept_lookup(name: str):
    """Look up a concept to find its table."""
    result = cache.lookup_concept(name)
    if result:
        click.echo(json.dumps(result, indent=2))
    else:
        click.echo(f"Concept '{name}' not found")


@concept.command("learn")
@click.argument("name")
@click.argument("table")
@click.option("--key-column", "-k", help="Primary key column")
@click.option("--date-column", "-d", help="Date column for filtering")
def concept_learn(name: str, table: str, key_column: str, date_column: str):
    """Store a concept -> table mapping."""
    cache.learn_concept(name, table, key_column, date_column)
    click.echo(f"Learned: '{name}' -> {table}")


@concept.command("list")
def concept_list():
    """List all learned concepts."""
    concepts = cache.list_concepts()
    if concepts:
        click.echo(json.dumps(concepts, indent=2))
    else:
        click.echo("No concepts cached yet")


@main.group()
def pattern():
    """Manage pattern cache (query strategies)."""
    pass


@pattern.command("lookup")
@click.argument("question")
def pattern_lookup(question: str):
    """Find patterns matching a question."""
    matches = cache.lookup_pattern(question)
    if matches:
        click.echo(json.dumps(matches, indent=2))
    else:
        click.echo("No matching patterns found")


@pattern.command("learn")
@click.argument("name")
@click.option(
    "--question-types",
    "-q",
    multiple=True,
    required=True,
    help="Question types this pattern handles",
)
@click.option("--strategy", "-s", multiple=True, required=True, help="Strategy steps")
@click.option("--tables", "-t", multiple=True, required=True, help="Tables used")
@click.option("--gotchas", "-g", multiple=True, help="Gotchas/warnings")
@click.option("--example", "-e", help="Example SQL query")
def pattern_learn(
    name: str,
    question_types: tuple,
    strategy: tuple,
    tables: tuple,
    gotchas: tuple,
    example: str,
):
    """Store a query pattern/strategy."""
    cache.learn_pattern(
        name=name,
        question_types=list(question_types),
        strategy=list(strategy),
        tables_used=list(tables),
        gotchas=list(gotchas),
        example_query=example,
    )
    click.echo(f"Learned pattern: '{name}'")


@pattern.command("record")
@click.argument("name")
@click.option("--success/--failure", default=True, help="Record success or failure")
def pattern_record(name: str, success: bool):
    """Record pattern outcome (success/failure)."""
    result = cache.record_pattern_outcome(name, success)
    if result:
        click.echo(f"Recorded {'success' if success else 'failure'} for '{name}'")
    else:
        click.echo(f"Pattern '{name}' not found")


@pattern.command("list")
def pattern_list():
    """List all learned patterns."""
    patterns = cache.list_patterns()
    if patterns:
        click.echo(json.dumps(patterns, indent=2))
    else:
        click.echo("No patterns cached yet")


@pattern.command("delete")
@click.argument("name")
def pattern_delete(name: str):
    """Delete a pattern by name."""
    if cache.delete_pattern(name):
        click.echo(f"Deleted pattern: '{name}'")
    else:
        click.echo(f"Pattern '{name}' not found")


# --- Cache Management ---


@main.group("cache")
def cache_group():
    """Manage cache (status, clear)."""
    pass


@cache_group.command("status")
def cache_status():
    """Show cache statistics."""
    stats = cache.cache_stats()
    click.echo(json.dumps(stats, indent=2))


@cache_group.command("clear")
@click.option(
    "--type",
    "cache_type",
    type=click.Choice(["all", "concepts", "patterns"]),
    default="all",
    help="What to clear",
)
@click.option("--stale-only", is_flag=True, help="Only clear entries older than TTL")
@click.confirmation_option(prompt="Are you sure you want to clear the cache?")
def cache_clear(cache_type: str, stale_only: bool):
    """Clear cache entries."""
    result = cache.clear_cache(cache_type, purge_stale_only=stale_only)
    click.echo(
        f"Cleared {result['concepts_cleared']} concepts, "
        f"{result['patterns_cleared']} patterns"
    )


# --- Table Schema Cache ---


@main.group()
def table():
    """Manage table schema cache."""
    pass


@table.command("lookup")
@click.argument("full_name")
def table_lookup(full_name: str):
    """Look up a cached table schema (DATABASE.SCHEMA.TABLE)."""
    result = cache.get_table(full_name)
    if result:
        click.echo(json.dumps(result, indent=2))
    else:
        click.echo(f"Table '{full_name}' not in cache")


@table.command("cache")
@click.argument("full_name")
@click.option("--columns", "-c", help="JSON array of column definitions")
@click.option("--row-count", "-r", type=int, help="Row count")
@click.option("--comment", help="Table description")
def table_cache(full_name: str, columns: str, row_count: int, comment: str):
    """Cache a table's schema.

    Example: uv run scripts/cli.py table cache DB.SCHEMA.TABLE -c '[{"name":"id","type":"INT"}]'
    """
    if columns:
        cols = json.loads(columns)
    else:
        cols = []
    cache.set_table(full_name, cols, row_count, comment)
    click.echo(f"Cached table: '{full_name}'")


@table.command("list")
def table_list():
    """List all cached table schemas."""
    tables = cache.list_tables()
    if tables:
        # Show summary (name + column count + cached_at)
        for name, info in tables.items():
            col_count = len(info.get("columns", []))
            cached_at = info.get("cached_at", "unknown")[:10]
            click.echo(f"{name}: {col_count} columns (cached {cached_at})")
    else:
        click.echo("No tables cached yet")


@table.command("delete")
@click.argument("full_name")
def table_delete(full_name: str):
    """Remove a table from cache."""
    if cache.delete_table(full_name):
        click.echo(f"Deleted table: '{full_name}'")
    else:
        click.echo(f"Table '{full_name}' not found")


# --- Bulk Import ---


@concept.command("import")
@click.option("--path", "-p", type=click.Path(exists=True), help="Path to warehouse.md")
def concept_import(path: str):
    """Import concepts from warehouse.md Quick Reference table.

    Parses markdown tables with: | Concept | Table | Key Column | Date Column |
    """
    from pathlib import Path as P

    file_path = P(path) if path else None
    count = cache.load_concepts_from_warehouse_md(file_path)
    if count > 0:
        click.echo(f"Imported {count} concepts from warehouse.md")
    else:
        click.echo("No concepts found in warehouse.md")


if __name__ == "__main__":
    main()

"""Configuration utilities for the analyzing-data skill."""

import sys
import warnings
from pathlib import Path

# Legacy path (deprecated)
_LEGACY_CONFIG_DIR = Path.home() / ".astro" / "ai" / "config"
# New path
_NEW_CONFIG_DIR = Path.home() / ".astro" / "agents"

_legacy_warning_shown = False


def _check_legacy_path() -> Path | None:
    """Check if legacy config path exists and warn user to migrate.

    Returns the legacy path if it exists and should be used, None otherwise.
    """
    global _legacy_warning_shown

    if _LEGACY_CONFIG_DIR.exists() and not _NEW_CONFIG_DIR.exists():
        if not _legacy_warning_shown:
            warnings.warn(
                f"Deprecated config path: {_LEGACY_CONFIG_DIR}\n"
                f"  Please move your config to: {_NEW_CONFIG_DIR}\n"
                f"  Run: mv ~/.astro/ai/config ~/.astro/agents",
                DeprecationWarning,
                stacklevel=3,
            )
            # Also print to stderr for CLI visibility
            print(
                "WARNING: Using deprecated config path ~/.astro/ai/config/\n"
                "  Please migrate: mv ~/.astro/ai/config ~/.astro/agents",
                file=sys.stderr,
            )
            _legacy_warning_shown = True
        return _LEGACY_CONFIG_DIR
    return None


def get_kernel_venv_dir() -> Path:
    """Get the path to the kernel virtual environment directory."""
    legacy = _check_legacy_path()
    if legacy:
        return legacy.parent / "kernel_venv"
    return _NEW_CONFIG_DIR / "kernel_venv"


def get_kernel_connection_file() -> Path:
    """Get the path to the kernel connection file."""
    legacy = _check_legacy_path()
    if legacy:
        return legacy.parent / "kernel.json"
    return _NEW_CONFIG_DIR / "kernel.json"


def get_config_dir() -> Path:
    """Get the path to the config directory."""
    legacy = _check_legacy_path()
    if legacy:
        return legacy
    return _NEW_CONFIG_DIR

"""Database connector registry, base class, and all connector implementations."""

import os
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, NamedTuple, TypeVar


# --- Base class ---


@dataclass
class DatabaseConnector(ABC):
    """Base class for database connectors."""

    databases: list[str]

    @classmethod
    @abstractmethod
    def connector_type(cls) -> str:
        """Return type identifier (e.g., 'snowflake', 'postgres')."""

    @classmethod
    @abstractmethod
    def from_dict(cls, data: dict[str, Any]) -> "DatabaseConnector":
        """Create from config dict."""

    @abstractmethod
    def validate(self, name: str) -> None:
        """Validate config. Raise ValueError if invalid."""

    @abstractmethod
    def get_required_packages(self) -> list[str]:
        """Return pip packages needed."""

    @abstractmethod
    def get_env_vars_for_kernel(self) -> dict[str, str]:
        """Return env vars to inject into kernel."""

    @abstractmethod
    def to_python_prelude(self) -> str:
        """Generate Python code for connection + helpers."""


# --- Utilities ---


def substitute_env_vars(value: Any) -> tuple[Any, str | None]:
    """Substitute ${VAR_NAME} with environment variable value."""
    if not isinstance(value, str):
        return value, None
    match = re.match(r"^\$\{([^}]+)\}$", value)
    if match:
        env_var_name = match.group(1)
        env_value = os.environ.get(env_var_name)
        return (env_value if env_value else value), env_var_name
    return value, None


# --- Registry ---

_CONNECTOR_REGISTRY: dict[str, type[DatabaseConnector]] = {}


T = TypeVar("T", bound="DatabaseConnector")


def register_connector(cls: type[T]) -> type[T]:
    _CONNECTOR_REGISTRY[cls.connector_type()] = cls
    return cls


def get_connector_class(connector_type: str) -> type[DatabaseConnector]:
    if connector_type not in _CONNECTOR_REGISTRY:
        available = ", ".join(sorted(_CONNECTOR_REGISTRY.keys()))
        raise ValueError(
            f"Unknown connector type: {connector_type!r}. Available: {available}"
        )
    return _CONNECTOR_REGISTRY[connector_type]


def create_connector(data: dict[str, Any]) -> DatabaseConnector:
    connector_type = data.get("type", "snowflake")
    cls = get_connector_class(connector_type)
    return cls.from_dict(data)


def list_connector_types() -> list[str]:
    return sorted(_CONNECTOR_REGISTRY.keys())


# --- Snowflake Connector ---


@register_connector
@dataclass
class SnowflakeConnector(DatabaseConnector):
    account: str = ""
    user: str = ""
    auth_type: str = "password"
    password: str = ""
    private_key_path: str = ""
    private_key_passphrase: str = ""
    private_key: str = ""
    warehouse: str = ""
    role: str = ""
    schema: str = ""
    databases: list[str] = field(default_factory=list)
    client_session_keep_alive: bool = False
    password_env_var: str | None = None
    private_key_env_var: str | None = None
    private_key_passphrase_env_var: str | None = None
    query_tag: str = ""

    @classmethod
    def connector_type(cls) -> str:
        return "snowflake"

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SnowflakeConnector":
        account, _ = substitute_env_vars(data.get("account", ""))
        user, _ = substitute_env_vars(data.get("user", ""))
        password, pw_env = substitute_env_vars(data.get("password", ""))
        private_key, pk_env = substitute_env_vars(data.get("private_key", ""))
        passphrase, pp_env = substitute_env_vars(data.get("private_key_passphrase", ""))

        return cls(
            account=account,
            user=user,
            auth_type=data.get("auth_type", "password"),
            password=password,
            private_key_path=data.get("private_key_path", ""),
            private_key_passphrase=passphrase,
            private_key=private_key,
            warehouse=data.get("warehouse", ""),
            role=data.get("role", ""),
            schema=data.get("schema", ""),
            databases=data.get("databases", []),
            client_session_keep_alive=data.get("client_session_keep_alive", False),
            password_env_var=pw_env,
            private_key_env_var=pk_env,
            private_key_passphrase_env_var=pp_env,
            query_tag=data.get("query_tag", ""),
        )

    def validate(self, name: str) -> None:
        if not self.account or self.account.startswith("${"):
            raise ValueError(f"warehouse '{name}': account required")
        if not self.user or self.user.startswith("${"):
            raise ValueError(f"warehouse '{name}': user required")
        if self.auth_type == "password":
            if not self.password or self.password.startswith("${"):
                raise ValueError(f"warehouse '{name}': password required")
        elif self.auth_type == "private_key":
            if not self.private_key_path and not self.private_key:
                raise ValueError(f"warehouse '{name}': private_key required")
        if len(self.query_tag) > 2000:
            raise ValueError(
                f"warehouse '{name}': query_tag exceeds Snowflake's 2000 character limit"
            )

    def get_required_packages(self) -> list[str]:
        pkgs = ["snowflake-connector-python[pandas]"]
        if self.auth_type == "private_key":
            pkgs.append("cryptography")
        return pkgs

    def get_env_vars_for_kernel(self) -> dict[str, str]:
        env_vars = {}
        if self.password_env_var and self.password:
            env_vars[self.password_env_var] = self.password
        if self.private_key_env_var and self.private_key:
            env_vars[self.private_key_env_var] = self.private_key
        if self.private_key_passphrase_env_var and self.private_key_passphrase:
            env_vars[self.private_key_passphrase_env_var] = self.private_key_passphrase
        return env_vars

    def to_python_prelude(self) -> str:
        from templates import (
            HELPERS_CODE,
            PRIVATE_KEY_CONTENT_TEMPLATE,
            PRIVATE_KEY_FILE_TEMPLATE,
        )

        sections = []

        # Imports
        sections.append("""import snowflake.connector
import polars as pl
import pandas as pd
import os""")

        # Private key loader (if needed)
        if self.auth_type == "private_key":
            if self.private_key_passphrase_env_var:
                passphrase_code = f"os.environ.get({self.private_key_passphrase_env_var!r}, '').encode() or None"
            elif self.private_key_passphrase:
                passphrase_code = f"{self.private_key_passphrase!r}.encode()"
            else:
                passphrase_code = "None"

            if self.private_key_path:
                sections.append(
                    PRIVATE_KEY_FILE_TEMPLATE.substitute(
                        KEY_PATH=repr(self.private_key_path),
                        PASSPHRASE_CODE=passphrase_code,
                    )
                )
            else:
                key_code = (
                    f"os.environ.get({self.private_key_env_var!r})"
                    if self.private_key_env_var
                    else repr(self.private_key)
                )
                sections.append(
                    PRIVATE_KEY_CONTENT_TEMPLATE.substitute(
                        KEY_CODE=key_code,
                        PASSPHRASE_CODE=passphrase_code,
                    )
                )

        # Connection
        lines = ["_conn = snowflake.connector.connect("]
        lines.append(f"    account={self.account!r},")
        lines.append(f"    user={self.user!r},")
        if self.auth_type == "password":
            if self.password_env_var:
                lines.append(f"    password=os.environ.get({self.password_env_var!r}),")
            else:
                lines.append(f"    password={self.password!r},")
        elif self.auth_type == "private_key":
            lines.append("    private_key=_load_private_key(),")
        if self.warehouse:
            lines.append(f"    warehouse={self.warehouse!r},")
        if self.role:
            lines.append(f"    role={self.role!r},")
        if self.databases:
            lines.append(f"    database={self.databases[0]!r},")
        if self.query_tag:
            lines.append(f"    session_parameters={{'QUERY_TAG': {self.query_tag!r}}},")
        lines.append(f"    client_session_keep_alive={self.client_session_keep_alive},")
        lines.append(")")
        sections.append("\n".join(lines))

        # Helper functions
        helpers_code = HELPERS_CODE
        if "def " in helpers_code:
            helpers_code = "def " + helpers_code.split("def ", 1)[1]
        sections.append(helpers_code.strip())

        # Status output
        status_lines = [
            'print("Snowflake connection established")',
            'print(f"   Account: {_conn.account}")',
            'print(f"   User: {_conn.user}")',
        ]
        if self.warehouse:
            status_lines.append(f'print(f"   Warehouse: {self.warehouse}")')
        if self.role:
            status_lines.append(f'print(f"   Role: {self.role}")')
        if self.databases:
            status_lines.append(f'print(f"   Database: {self.databases[0]}")')
        if self.query_tag:
            status_lines.append(f'print(f"   Query Tag: {self.query_tag}")')
        status_lines.append(
            'print("\\nAvailable: run_sql(query) -> polars, run_sql_pandas(query) -> pandas")'
        )
        sections.append("\n".join(status_lines))

        return "\n\n".join(sections)


# --- PostgreSQL Connector ---


@register_connector
@dataclass
class PostgresConnector(DatabaseConnector):
    host: str = ""
    port: int = 5432
    user: str = ""
    password: str = ""
    database: str = ""
    sslmode: str = ""
    databases: list[str] = field(default_factory=list)
    password_env_var: str | None = None
    application_name: str = ""

    @classmethod
    def connector_type(cls) -> str:
        return "postgres"

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "PostgresConnector":
        host, _ = substitute_env_vars(data.get("host", ""))
        user, _ = substitute_env_vars(data.get("user", ""))
        password, pw_env = substitute_env_vars(data.get("password", ""))
        database, _ = substitute_env_vars(data.get("database", ""))

        return cls(
            host=host,
            port=data.get("port", 5432),
            user=user,
            password=password,
            database=database,
            sslmode=data.get("sslmode", ""),
            databases=data.get("databases", [database] if database else []),
            password_env_var=pw_env,
            application_name=data.get("application_name", ""),
        )

    def validate(self, name: str) -> None:
        if not self.host or self.host.startswith("${"):
            raise ValueError(f"warehouse '{name}': host required for postgres")
        if not self.user or self.user.startswith("${"):
            raise ValueError(f"warehouse '{name}': user required for postgres")
        if not self.database or self.database.startswith("${"):
            raise ValueError(f"warehouse '{name}': database required for postgres")

    def get_required_packages(self) -> list[str]:
        return ["psycopg[binary,pool]"]

    def get_env_vars_for_kernel(self) -> dict[str, str]:
        env_vars = {}
        if self.password_env_var and self.password:
            env_vars[self.password_env_var] = self.password
        return env_vars

    def to_python_prelude(self) -> str:
        lines = ["_conn = psycopg.connect("]
        lines.append(f"    host={self.host!r},")
        lines.append(f"    port={self.port},")
        lines.append(f"    user={self.user!r},")
        if self.password_env_var:
            lines.append(f"    password=os.environ.get({self.password_env_var!r}),")
        elif self.password:
            lines.append(f"    password={self.password!r},")
        lines.append(f"    dbname={self.database!r},")
        if self.sslmode:
            lines.append(f"    sslmode={self.sslmode!r},")
        if self.application_name:
            lines.append(f"    application_name={self.application_name!r},")
        lines.append("    autocommit=True,")
        lines.append(")")
        connection_code = "\n".join(lines)

        status_lines = [
            'print("PostgreSQL connection established")',
            f'print("   Host: {self.host}:{self.port}")',
            f'print("   User: {self.user}")',
            f'print("   Database: {self.database}")',
        ]
        if self.application_name:
            status_lines.append(f'print("   Application: {self.application_name}")')
        status_lines += [
            'print("\\nAvailable: run_sql(query) -> polars, run_sql_pandas(query) -> pandas")',
        ]
        status_code = "\n".join(status_lines)

        return f'''import psycopg
import polars as pl
import pandas as pd
import os

{connection_code}

def run_sql(query: str, limit: int = 100):
    """Execute SQL and return Polars DataFrame."""
    with _conn.cursor() as cursor:
        cursor.execute(query)
        if cursor.description is None:
            return pl.DataFrame()
        columns = [desc[0] for desc in cursor.description]
        rows = cursor.fetchall()
        result = pl.DataFrame(rows, schema=columns, orient="row")
        return result.head(limit) if limit > 0 and len(result) > limit else result


def run_sql_pandas(query: str, limit: int = 100):
    """Execute SQL and return Pandas DataFrame."""
    with _conn.cursor() as cursor:
        cursor.execute(query)
        if cursor.description is None:
            return pd.DataFrame()
        columns = [desc[0] for desc in cursor.description]
        rows = cursor.fetchall()
        df = pd.DataFrame(rows, columns=columns)
        return df.head(limit) if limit > 0 and len(df) > limit else df

{status_code}'''


# --- BigQuery Connector ---

# Google allows international characters in BQ labels, but we restrict to ASCII
# for simplicity. Expand the regex if international support is needed.
_BQ_LABEL_KEY_RE = re.compile(r"^[a-z][a-z0-9_-]{0,62}$")
_BQ_LABEL_VALUE_RE = re.compile(r"^[a-z0-9_-]{0,63}$")


@register_connector
@dataclass
class BigQueryConnector(DatabaseConnector):
    project: str = ""
    credentials_path: str = ""
    location: str = ""
    databases: list[str] = field(default_factory=list)
    labels: dict[str, str] = field(default_factory=dict)

    @classmethod
    def connector_type(cls) -> str:
        return "bigquery"

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "BigQueryConnector":
        project, _ = substitute_env_vars(data.get("project", ""))
        credentials_path, _ = substitute_env_vars(data.get("credentials_path", ""))

        return cls(
            project=project,
            credentials_path=credentials_path,
            location=data.get("location", ""),
            databases=data.get("databases", [project] if project else []),
            labels=data.get("labels", {}),
        )

    def validate(self, name: str) -> None:
        if not self.project or self.project.startswith("${"):
            raise ValueError(f"warehouse '{name}': project required for bigquery")
        if len(self.labels) > 64:
            raise ValueError(
                f"warehouse '{name}': BigQuery supports at most 64 labels, got {len(self.labels)}"
            )
        for k, v in self.labels.items():
            if not isinstance(k, str) or not _BQ_LABEL_KEY_RE.match(k):
                raise ValueError(
                    f"warehouse '{name}': invalid BigQuery label key {k!r} "
                    "(must match [a-z][a-z0-9_-]{0,62})"
                )
            if not isinstance(v, str):
                raise ValueError(
                    f"warehouse '{name}': label value for {k!r} must be a string, got {type(v).__name__}"
                )
            if not _BQ_LABEL_VALUE_RE.match(v):
                raise ValueError(
                    f"warehouse '{name}': invalid BigQuery label value {v!r} for key {k!r} "
                    "(must match [a-z0-9_-]{0,63})"
                )

    def get_required_packages(self) -> list[str]:
        return ["google-cloud-bigquery[pandas,pyarrow]", "db-dtypes"]

    def get_env_vars_for_kernel(self) -> dict[str, str]:
        env_vars = {}
        if self.credentials_path:
            env_vars["GOOGLE_APPLICATION_CREDENTIALS"] = self.credentials_path
        return env_vars

    def to_python_prelude(self) -> str:
        if self.credentials_path:
            conn_code = f"""from google.oauth2 import service_account
_credentials = service_account.Credentials.from_service_account_file({self.credentials_path!r})
_client = bigquery.Client(project={self.project!r}, credentials=_credentials)"""
        else:
            conn_code = f"_client = bigquery.Client(project={self.project!r})"

        # Build QueryJobConfig arguments
        job_config_args = []
        if self.labels:
            job_config_args.append(f"labels={self.labels!r}")
        job_config_str = ", ".join(job_config_args)

        # Build _client.query() extra kwargs
        query_extra_args = ""
        if self.location:
            query_extra_args = f", location={self.location!r}"

        auth_type = (
            "Service Account"
            if self.credentials_path
            else "Application Default Credentials"
        )

        status_lines = [
            'print("BigQuery client initialized")',
            f'print(f"   Project: {self.project}")',
        ]
        if self.location:
            status_lines.append(f'print(f"   Location: {self.location}")')
        status_lines.append(f'print("   Auth: {auth_type}")')
        if self.labels:
            status_lines.append(f'print(f"   Labels: {self.labels!r}")')
        status_lines.append(
            'print("\\nAvailable: run_sql(query) -> polars, run_sql_pandas(query) -> pandas")'
        )
        status_code = "\n".join(status_lines)

        return f'''from google.cloud import bigquery
import polars as pl
import pandas as pd
import os

{conn_code}

def run_sql(query: str, limit: int = 100):
    """Execute SQL and return Polars DataFrame."""
    job_config = bigquery.QueryJobConfig({job_config_str})
    query_job = _client.query(query, job_config=job_config{query_extra_args})
    df = query_job.to_dataframe()
    result = pl.from_pandas(df)
    return result.head(limit) if limit > 0 and len(result) > limit else result


def run_sql_pandas(query: str, limit: int = 100):
    """Execute SQL and return Pandas DataFrame."""
    job_config = bigquery.QueryJobConfig({job_config_str})
    query_job = _client.query(query, job_config=job_config{query_extra_args})
    df = query_job.to_dataframe()
    return df.head(limit) if limit > 0 and len(df) > limit else df

{status_code}'''


# --- SQLAlchemy Connector ---


class DialectInfo(NamedTuple):
    """Database dialect configuration.

    To add a new database:
    1. Add an entry to DIALECTS below with (display_name, [packages])
    2. Run tests: uv run pytest tests/test_connectors.py -v
    """

    display_name: str
    packages: list[str]


# Mapping of dialect/driver names to their configuration.
# The dialect is extracted from URLs like "dialect+driver://..." or "dialect://..."
# When a driver is specified (e.g., mysql+pymysql), the driver name is looked up first.
DIALECTS: dict[str, DialectInfo] = {
    # PostgreSQL variants
    "postgresql": DialectInfo("PostgreSQL", ["psycopg[binary]"]),
    "postgres": DialectInfo("PostgreSQL", ["psycopg[binary]"]),
    "psycopg": DialectInfo("PostgreSQL", ["psycopg[binary]"]),
    "psycopg2": DialectInfo("PostgreSQL", ["psycopg2-binary"]),
    "pg8000": DialectInfo("PostgreSQL", ["pg8000"]),
    "asyncpg": DialectInfo("PostgreSQL", ["asyncpg"]),
    # MySQL variants
    "mysql": DialectInfo("MySQL", ["pymysql"]),
    "pymysql": DialectInfo("MySQL", ["pymysql"]),
    "mysqlconnector": DialectInfo("MySQL", ["mysql-connector-python"]),
    "mysqldb": DialectInfo("MySQL", ["mysqlclient"]),
    "mariadb": DialectInfo("MariaDB", ["mariadb"]),
    # SQLite (built-in, no extra packages)
    "sqlite": DialectInfo("SQLite", []),
    # Oracle
    "oracle": DialectInfo("Oracle", ["oracledb"]),
    "oracledb": DialectInfo("Oracle", ["oracledb"]),
    # SQL Server
    "mssql": DialectInfo("SQL Server", ["pyodbc"]),
    "pyodbc": DialectInfo("SQL Server", ["pyodbc"]),
    "pymssql": DialectInfo("SQL Server", ["pymssql"]),
    # Cloud data warehouses
    "redshift": DialectInfo("Redshift", ["redshift_connector"]),
    "redshift_connector": DialectInfo("Redshift", ["redshift_connector"]),
    "snowflake": DialectInfo(
        "Snowflake", ["snowflake-sqlalchemy", "snowflake-connector-python"]
    ),
    "bigquery": DialectInfo("BigQuery", ["sqlalchemy-bigquery"]),
    # DuckDB
    "duckdb": DialectInfo("DuckDB", ["duckdb", "duckdb-engine"]),
    # Other databases
    "trino": DialectInfo("Trino", ["trino"]),
    "clickhouse": DialectInfo(
        "ClickHouse", ["clickhouse-driver", "clickhouse-sqlalchemy"]
    ),
    "cockroachdb": DialectInfo(
        "CockroachDB", ["sqlalchemy-cockroachdb", "psycopg[binary]"]
    ),
    "databricks": DialectInfo("Databricks", ["databricks-sql-connector"]),
    "teradata": DialectInfo("Teradata", ["teradatasqlalchemy"]),
    "vertica": DialectInfo("Vertica", ["vertica-python"]),
    "hana": DialectInfo("SAP HANA", ["hdbcli"]),
    "db2": DialectInfo("IBM Db2", ["ibm_db_sa"]),
    "firebird": DialectInfo("Firebird", ["fdb"]),
    "awsathena": DialectInfo("Amazon Athena", ["pyathena"]),
    "spanner": DialectInfo("Cloud Spanner", ["sqlalchemy-spanner"]),
}


def _extract_dialect(url: str) -> str | None:
    """Extract dialect name from SQLAlchemy URL.

    URLs can be:
    - dialect://user:pass@host/db
    - dialect+driver://user:pass@host/db

    When a driver is specified, returns the driver name (looked up first in DIALECTS).
    Falls back to dialect name if driver isn't in DIALECTS.
    """
    match = re.match(r"^([a-zA-Z0-9_-]+)(?:\+([a-zA-Z0-9_-]+))?://", url)
    if match:
        dialect = match.group(1).lower()
        driver = match.group(2).lower() if match.group(2) else None
        # Prefer driver if specified AND it's in our dialects mapping
        # Otherwise fall back to dialect (e.g., postgresql+asyncpg -> asyncpg if known)
        if driver and driver in DIALECTS:
            return driver
        return dialect
    return None


@register_connector
@dataclass
class SQLAlchemyConnector(DatabaseConnector):
    url: str = ""
    databases: list[str] = field(default_factory=list)
    pool_size: int = 5
    echo: bool = False
    url_env_var: str | None = None
    connect_args: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def connector_type(cls) -> str:
        return "sqlalchemy"

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SQLAlchemyConnector":
        url, url_env = substitute_env_vars(data.get("url", ""))

        return cls(
            url=url,
            databases=data.get("databases", []),
            pool_size=data.get("pool_size", 5),
            echo=data.get("echo", False),
            url_env_var=url_env,
            connect_args=data.get("connect_args", {}),
        )

    def validate(self, name: str) -> None:
        if not self.url or self.url.startswith("${"):
            raise ValueError(f"warehouse '{name}': url required for sqlalchemy")
        if not self.databases:
            raise ValueError(
                f"warehouse '{name}': databases list required for sqlalchemy"
            )

    def get_required_packages(self) -> list[str]:
        packages = ["sqlalchemy"]
        dialect = _extract_dialect(self.url)
        if dialect and dialect in DIALECTS:
            packages.extend(DIALECTS[dialect].packages)
        return packages

    def get_env_vars_for_kernel(self) -> dict[str, str]:
        env_vars = {}
        if self.url_env_var and self.url:
            env_vars[self.url_env_var] = self.url
        return env_vars

    def to_python_prelude(self) -> str:
        if self.url_env_var:
            url_code = f"os.environ.get({self.url_env_var!r})"
        else:
            url_code = repr(self.url)

        # Infer DB type for status message
        dialect = _extract_dialect(self.url)
        db_type = (
            DIALECTS[dialect].display_name
            if dialect and dialect in DIALECTS
            else "Database"
        )

        databases_str = ", ".join(self.databases)

        return f'''from sqlalchemy import create_engine, text
import polars as pl
import pandas as pd
import os
import atexit

_engine = create_engine({url_code}, pool_size={self.pool_size}, echo={self.echo}{f", connect_args={self.connect_args!r}" if self.connect_args else ""})
_conn = _engine.connect()
atexit.register(lambda: (_conn.close(), _engine.dispose()))

def run_sql(query: str, limit: int = 100):
    """Execute SQL and return Polars DataFrame."""
    result = _conn.execute(text(query))
    if result.returns_rows:
        columns = list(result.keys())
        rows = result.fetchall()
        df = pl.DataFrame(rows, schema=columns, orient="row")
        return df.head(limit) if limit > 0 and len(df) > limit else df
    return pl.DataFrame()


def run_sql_pandas(query: str, limit: int = 100):
    """Execute SQL and return Pandas DataFrame."""
    result = _conn.execute(text(query))
    if result.returns_rows:
        columns = list(result.keys())
        rows = result.fetchall()
        df = pd.DataFrame(rows, columns=columns)
        return df.head(limit) if limit > 0 and len(df) > limit else df
    return pd.DataFrame()

print("{db_type} connection established (via SQLAlchemy)")
print(f"   Database(s): {databases_str}")
print("\\nAvailable: run_sql(query) -> polars, run_sql_pandas(query) -> pandas")'''


__all__ = [
    "DatabaseConnector",
    "substitute_env_vars",
    "register_connector",
    "get_connector_class",
    "create_connector",
    "list_connector_types",
    "SnowflakeConnector",
    "PostgresConnector",
    "BigQueryConnector",
    "SQLAlchemyConnector",
    "DialectInfo",
    "DIALECTS",
]

"""Jupyter kernel manager for executing Python code with persistent state."""

import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

from jupyter_client import KernelManager as JupyterKernelManager
from jupyter_client import BlockingKernelClient

from config import get_kernel_venv_dir, get_kernel_connection_file

DEFAULT_PACKAGES = [
    "ipykernel",
    "jupyter_client",
    "polars",
    "pandas",
    "numpy",
    "matplotlib",
    "seaborn",
    "pyyaml",
    "python-dotenv",
]


@dataclass
class ExecutionResult:
    """Result of code execution in the kernel."""

    success: bool
    output: str
    error: str | None = None


class KernelManager:
    """Manages a Jupyter kernel for Python code execution."""

    def __init__(
        self,
        venv_dir: Path | None = None,
        kernel_name: str = "astro-ai-kernel",
        packages: list[str] | None = None,
    ):
        self.venv_dir = venv_dir or get_kernel_venv_dir()
        self.kernel_name = kernel_name
        self.packages = packages or DEFAULT_PACKAGES.copy()
        self.connection_file = get_kernel_connection_file()
        self._km: JupyterKernelManager | None = None

    @property
    def python_path(self) -> Path:
        if sys.platform == "win32":
            return self.venv_dir / "Scripts" / "python.exe"
        return self.venv_dir / "bin" / "python"

    @property
    def is_running(self) -> bool:
        if not self.connection_file.exists():
            return False
        try:
            kc = BlockingKernelClient()
            kc.load_connection_file(str(self.connection_file))
            kc.start_channels()
            try:
                kc.wait_for_ready(timeout=2)
                return True
            except Exception:
                return False
            finally:
                kc.stop_channels()
        except Exception:
            return False

    def ensure_environment(self, extra_packages: list[str] | None = None) -> None:
        if not shutil.which("uv"):
            raise RuntimeError(
                "uv is not installed.\n"
                "Install with: curl -LsSf https://astral.sh/uv/install.sh | sh"
            )

        packages = self.packages.copy()
        if extra_packages:
            packages.extend(extra_packages)

        if not self.venv_dir.exists():
            print(f"Creating environment at {self.venv_dir}")
            subprocess.run(
                ["uv", "venv", str(self.venv_dir), "--seed"],
                check=True,
                capture_output=True,
            )

        print("Installing packages...")
        subprocess.run(
            ["uv", "pip", "install", "--python", str(self.python_path)] + packages,
            check=True,
            capture_output=True,
        )

        # Register kernel
        try:
            subprocess.run(
                [
                    str(self.python_path),
                    "-m",
                    "ipykernel",
                    "install",
                    "--user",
                    "--name",
                    self.kernel_name,
                    "--display-name",
                    "Data Analysis Kernel",
                ],
                capture_output=True,
                timeout=30,
            )
        except Exception:
            pass

    def start(
        self,
        env_vars: dict[str, str] | None = None,
        extra_packages: list[str] | None = None,
    ) -> None:
        if self.is_running:
            print("Kernel already running")
            return

        self.ensure_environment(extra_packages=extra_packages)
        print("Starting kernel...")

        self._km = JupyterKernelManager(kernel_name=self.kernel_name)

        if env_vars:
            import os

            for key, value in env_vars.items():
                os.environ[key] = value

        self._km.start_kernel(extra_arguments=["--IPKernelApp.parent_handle=0"])

        self.connection_file.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(self._km.connection_file, self.connection_file)

        kc = self._km.client()
        kc.start_channels()
        try:
            kc.wait_for_ready(timeout=10)
        except Exception as e:
            self.stop()
            raise RuntimeError(f"Kernel failed: {e}") from e
        finally:
            kc.stop_channels()

        # Inject idle timeout watchdog into the kernel
        self._km.client().execute(
            "import threading, time, os, signal\n"
            "_idle_timeout = 1800\n"  # 30 minutes
            "_last_active = [time.time()]\n"
            "_orig_execute = get_ipython().run_cell\n"
            "def _tracked_execute(*a, **kw):\n"
            "    _last_active[0] = time.time()\n"
            "    return _orig_execute(*a, **kw)\n"
            "get_ipython().run_cell = _tracked_execute\n"
            "def _idle_watchdog():\n"
            "    while True:\n"
            "        time.sleep(60)\n"
            "        if time.time() - _last_active[0] > _idle_timeout:\n"
            "            os._exit(0)\n"
            "_t = threading.Thread(target=_idle_watchdog, daemon=True)\n"
            "_t.start()\n",
            silent=True,
        )

        self._km = None
        print(f"Kernel started ({self.connection_file})")

    def stop(self) -> None:
        if not self.connection_file.exists():
            print("Kernel not running")
            return

        try:
            kc = BlockingKernelClient()
            kc.load_connection_file(str(self.connection_file))
            kc.start_channels()
            kc.shutdown()
            kc.stop_channels()
        except Exception:
            pass

        if self.connection_file.exists():
            self.connection_file.unlink()
        print('{"message": "Kernel stopped"}')

    def execute(self, code: str, timeout: float = 30.0) -> ExecutionResult:
        if not self.connection_file.exists():
            return ExecutionResult(
                False, "", "Kernel not running. Start with: uv run scripts/cli.py start"
            )

        kc = BlockingKernelClient()
        kc.load_connection_file(str(self.connection_file))
        kc.start_channels()

        try:
            kc.wait_for_ready(timeout=5)
        except Exception as e:
            kc.stop_channels()
            return ExecutionResult(False, "", f"Kernel not responding: {e}")

        msg_id = kc.execute(code, silent=False, store_history=True)

        output_parts: list[str] = []
        error_msg: str | None = None
        status = "ok"
        deadline = time.time() + timeout
        done = False

        while time.time() < deadline and not done:
            try:
                msg = kc.get_iopub_msg(timeout=min(1.0, deadline - time.time()))
                if msg["parent_header"].get("msg_id") != msg_id:
                    continue

                msg_type = msg["msg_type"]
                content = msg["content"]

                if msg_type == "stream":
                    output_parts.append(content["text"])
                elif msg_type == "execute_result":
                    output_parts.append(content["data"].get("text/plain", ""))
                elif msg_type == "error":
                    error_msg = "\n".join(content["traceback"])
                    status = "error"
                elif msg_type == "status" and content["execution_state"] == "idle":
                    done = True
            except Exception:
                continue

        kc.stop_channels()

        if not done:
            return ExecutionResult(
                False, "".join(output_parts), f"Timeout after {timeout}s"
            )

        return ExecutionResult(status == "ok", "".join(output_parts), error_msg)

    def status(self) -> dict:
        info = {
            "running": False,
            "connection_file": str(self.connection_file),
            "responsive": False,
        }
        if not self.connection_file.exists():
            return info
        info["running"] = True
        try:
            kc = BlockingKernelClient()
            kc.load_connection_file(str(self.connection_file))
            kc.start_channels()
            try:
                kc.wait_for_ready(timeout=2)
                info["responsive"] = True
            except Exception:
                pass
            finally:
                kc.stop_channels()
        except Exception:
            pass
        return info

    def install_packages(self, packages: list[str]) -> tuple[bool, str]:
        """Install additional packages into the kernel environment.

        Args:
            packages: List of package specs (e.g., ['plotly>=5.0', 'scipy'])

        Returns:
            Tuple of (success, message)
        """
        if not packages:
            return False, "No packages specified"

        if not shutil.which("uv"):
            return False, "uv is not installed"

        try:
            result = subprocess.run(
                ["uv", "pip", "install", "--python", str(self.python_path)] + packages,
                capture_output=True,
                text=True,
            )
            if result.returncode == 0:
                return True, f"Installed: {', '.join(packages)}"
            else:
                return False, f"Failed: {result.stderr}"
        except Exception as e:
            return False, f"Error: {e}"