import typer
import uvicorn
import os
import sys
import importlib.util
from rich.console import Console
from rich.panel import Panel
from rich.layout import Layout
from rich.table import Table
from datetime import datetime
from .core import FeatureStore
from .server import create_app


app = typer.Typer()
console = Console()


@app.command(name="events")
def events_cmd(
    action: str = typer.Argument(..., help="Action: listen"),
    stream: str = typer.Option("meridian_events", help="Stream key to listen to"),
    count: int = typer.Option(10, help="Number of events to fetch per poll"),
    redis_url: str = typer.Option(
        None, envvar="MERIDIAN_REDIS_URL", help="Redis URL Override"
    ),
) -> None:
    """
    Manage or listen to Meridian events.
    Usage: meridian events listen --stream=my_stream
    """
    if action != "listen":
        console.print(f"[bold red]Unknown action:[/bold red] {action}")
        raise typer.Exit(1)

    import asyncio
    from redis.asyncio import Redis

    url = redis_url or os.getenv("MERIDIAN_REDIS_URL") or "redis://localhost:6379"

    async def listen_loop() -> None:
        console.print(f"[green]Listening to stream:[/green] {stream} on {url}")
        r = Redis.from_url(url, decode_responses=True)
        last_id = "$"

        try:
            while True:
                # XREAD block=0 means block indefinitely until new item
                streams = await r.xread({stream: last_id}, count=1, block=0)
                if not streams:
                    continue

                for stream_name, messages in streams:
                    for msg_id, fields in messages:
                        timestamp = datetime.now().strftime("%H:%M:%S")
                        console.print(
                            f"[{timestamp}] [bold cyan]{msg_id}[/bold cyan]: {fields}"
                        )
                        last_id = msg_id
        except asyncio.CancelledError:
            pass
        except Exception as e:
            console.print(f"[bold red]Error:[/bold red] {e}")
        finally:
            await r.aclose()  # type: ignore[attr-defined]

    try:
        asyncio.run(listen_loop())
    except KeyboardInterrupt:
        console.print("\nStopped.")


console = Console()


@app.callback()
def callback() -> None:
    """
    Meridian CLI
    """


@app.command(name="worker")
def worker_cmd(
    file: str = typer.Argument(
        ..., help="Path to the feature definition file (e.g., features.py)"
    ),
    redis_url: str = typer.Option(
        None, envvar="MERIDIAN_REDIS_URL", help="Redis URL Override"
    ),
) -> None:
    """
    Starts the Axiom background worker.
    """
    import asyncio
    from .worker import AxiomWorker

    # 1. Load Feature Definitions
    if not os.path.exists(file):
        console.print(f"[bold red]Error:[/bold red] File '{file}' not found.")
        raise typer.Exit(code=1)

    sys.path.append(os.getcwd())
    try:
        module_name = os.path.splitext(os.path.basename(file))[0]
        spec = importlib.util.spec_from_file_location(module_name, file)
        if spec and spec.loader:
            module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = module
            spec.loader.exec_module(module)

        # Find store instance
        store = None
        for attr_name in dir(module):
            attr = getattr(module, attr_name)
            if isinstance(attr, FeatureStore):
                store = attr
                break

        if not store:
            console.print("[bold red]Error:[/bold red] No FeatureStore found in file.")
            raise typer.Exit(code=1)

    except Exception as e:
        console.print(f"[bold red]Error loading features:[/bold red] {e}")
        raise typer.Exit(code=1)

    console.print(Panel(f"Starting Axiom Worker for {file}...", style="bold green"))

    # Pass store to worker
    worker = AxiomWorker(redis_url=redis_url, store=store)
    try:
        asyncio.run(worker.run())
    except KeyboardInterrupt:
        console.print("Worker stopped.")


@app.command(name="setup")
def setup(
    dir: str = typer.Argument(".", help="Directory to create setup files in"),
) -> None:
    """
    Generate production-ready configuration files (Docker Compose).
    """
    # Ensure directory exists
    if not os.path.exists(dir):
        os.makedirs(dir)
        console.print(f"Created directory: [bold cyan]{dir}[/bold cyan]")
    docker_compose = """
version: '3.8'

services:
  # 1. Postgres with pgvector (Offline Store + Context Store)
  postgres:
    image: pgvector/pgvector:pg16
    ports:
      - "5432:5432"
    environment:
      POSTGRES_USER: user
      POSTGRES_PASSWORD: password
      POSTGRES_DB: meridian
    volumes:
      - meridian_postgres_data:/var/lib/postgresql/data

  # 2. Redis (Online Store + Cache)
  redis:
    image: redis:alpine
    ports:
      - "6379:6379"
    volumes:
      - meridian_redis_data:/data

volumes:
  meridian_postgres_data:
  meridian_redis_data:
"""
    env_example = """
# Meridian Production Config

# Security
MERIDIAN_API_KEY=change_me_to_something_secure

# Data Stores
MERIDIAN_REDIS_URL=redis://localhost:6379
MERIDIAN_POSTGRES_URL=postgresql://user:password@localhost:5432/meridian  # pragma: allowlist secret

# LLM Providers (Required for Context Store)
OPENAI_API_KEY=sk-...
ANTHROPIC_API_KEY=sk-ant-...
COHERE_API_KEY=...
"""

    # Write files
    dc_path = os.path.join(dir, "docker-compose.yml")
    env_path = os.path.join(dir, ".env.production")

    if os.path.exists(dc_path):
        console.print(f"[yellow]Warning:[/yellow] {dc_path} already exists. Skipping.")
    else:
        with open(dc_path, "w") as f:
            f.write(docker_compose.strip())
        console.print("Created [bold]docker-compose.yml[/bold]")

    if os.path.exists(env_path):
        console.print(f"[yellow]Warning:[/yellow] {env_path} already exists. Skipping.")
    else:
        with open(env_path, "w") as f:
            f.write(env_example.strip())
        os.chmod(env_path, 0o600)
        console.print(
            "Created [bold].env.production[/bold] (permissions restricted to 0600)"
        )

    console.print("\n[green]Setup Complete![/green]")
    console.print("To start infrastructure, run:")
    console.print("  [bold]docker compose up -d[/bold]")


@app.command(name="init")
def init(
    name: str = typer.Argument("meridian_project", help="Project name"),
    demo: bool = typer.Option(False, help="Include demo features and data"),
    interactive: bool = typer.Option(True, help="Run in interactive mode"),
) -> None:
    """
    Initialize a new Meridian project.
    """
    if os.path.exists(name):
        console.print(f"[bold red]Error:[/bold red] Directory '{name}' already exists.")
        raise typer.Exit(1)

    os.makedirs(name)
    console.print(f"Created directory: [bold cyan]{name}[/bold cyan]")

    # Interactive Configuration
    api_key_lines = []
    if interactive:
        # Check if we are in a TTY
        if sys.stdin.isatty():
            console.print("\n[bold]Configuration[/bold]")
            target_provider = typer.prompt(
                "Which LLM provider do you want to configure? (openai/anthropic/cohere/skip)",
                default="skip",
            ).lower()

            if target_provider == "openai":
                k = typer.prompt("Enter OpenAI API Key", hide_input=True)
                api_key_lines.append(f"OPENAI_API_KEY={k}")
            elif target_provider == "anthropic":
                k = typer.prompt("Enter Anthropic API Key", hide_input=True)
                api_key_lines.append(f"ANTHROPIC_API_KEY={k}")
            elif target_provider == "cohere":
                k = typer.prompt("Enter Cohere API Key", hide_input=True)
                api_key_lines.append(f"COHERE_API_KEY={k}")

    # Basic scaffold
    gitignore = """
__pycache__/
*.pyc
.env
.venv
*.db
    """
    with open(os.path.join(name, ".gitignore"), "w") as f:
        f.write(gitignore.strip())

    if api_key_lines:
        env_path = os.path.join(name, ".env")
        with open(env_path, "w") as f:
            f.write("\n".join(api_key_lines) + "\n")
        os.chmod(env_path, 0o600)
        console.print(
            "Created [bold].env[/bold] with API key (permissions restricted to 0600)."
        )

    if demo:
        # Create features.py
        features_py = """
from meridian.core import FeatureStore, entity, feature
from meridian.context import context, Context, ContextItem
from meridian.retrieval import retriever
import random
import os

# Use default local stack (DuckDB + In-Memory)
store = FeatureStore()

@entity(store)
class User:
    user_id: str

# 1. Standard ML Feature
@feature(entity=User, refresh="5m", materialize=True)
def engagement_score(user_id: str) -> float:
    # Simulate a score
    return round(random.random() * 100, 2)

# 2. RAG / Context Store
# We assume there is an index called 'knowledge_base' (created via seed.py)

@retriever(store, index="knowledge_base", top_k=2)
async def semantic_search(query: str) -> list[str]:
    # In a real app with pgvector, this searches vectors.
    # For this local demo without Postgres/OpenAI keys, we mock the return
    # if the index isn't reachable or keys aren't set.
    return [
        "Meridian allows defining features in Python.",
        "The Context Store manages token budgets for LLMs."
    ]

@context(store, max_tokens=1000)
async def chatbot_context(user_id: str, query: str) -> Context:
    # Fetch data in parallel
    score = await store.get_feature("engagement_score", user_id)
    docs = await semantic_search(query)

    return Context(items=[
        ContextItem(f"User Engagement: {score}", priority=2),
        ContextItem(docs, priority=1, required=True),
        ContextItem("System: You are a helpful assistant.", priority=0, required=True),
    ])
"""
        with open(os.path.join(name, "features.py"), "w") as f:
            f.write(features_py.strip())

        # Create README
        readme = """
# Meridian Demo Project

This is a generated demo project.

## Quickstart

1. **Install Meridian**:
   ```bash
   pip install "meridian-oss[ui]"
   ```

2. **Run the Server**:
   ```bash
   meridian serve features.py
   ```

3. **Query Context (E.g. for User 'u1')**:
   ```bash
   meridian context explain u1 --query "What is Meridian?"
   ```
"""
        with open(os.path.join(name, "README.md"), "w") as f:
            f.write(readme.strip())

        console.print(f"[green]Initialized demo project in '{name}'[/green]")
        console.print(
            "Run [bold]meridian serve features.py[/bold] inside the directory to start."
        )

    else:
        # Empty features.py
        with open(os.path.join(name, "features.py"), "w") as f:
            f.write(
                "from meridian.core import FeatureStore\n\nstore = FeatureStore()\n"
            )
        console.print(f"[green]Initialized empty project in '{name}'[/green]")


@app.command(name="version")
def version() -> None:
    """
    Prints the Meridian version.
    """
    try:
        from importlib.metadata import version

        v = version("meridian-oss")
    except Exception:
        v = "unknown"
    console.print(f"Meridian OSS v{v}")


@app.command(name="serve")
def serve(
    file: str = typer.Argument(
        ..., help="Path to the feature definition file (e.g., features.py)"
    ),
    host: str = typer.Option("127.0.0.1", help="Host to bind to"),
    port: int = typer.Option(8000, help="Port to bind to"),
    api_key: str = typer.Option(
        None, envvar="MERIDIAN_API_KEY", help="API Key for security"
    ),
    reload: bool = typer.Option(False, help="Enable auto-reload"),
) -> None:
    """
    Starts the Meridian server.
    """
    if not os.path.exists(file):
        console.print(f"[bold red]Error:[/bold red] File '{file}' not found.")
        raise typer.Exit(code=1)

    console.print(
        Panel(
            f"Starting Meridian on http://{host}:{port}",
            title="Meridian",
            style="bold blue",
        )
    )

    # For now, we just print that we are starting.
    # In Epic 4 (Serving API), we will actually start the FastAPI server here.
    # For Epic 2, we just need to verify the CLI works and can load the file.

    # Simulate loading the file to register features
    sys.path.append(os.getcwd())

    try:
        # Import the module to execute the decorators and register features
        module_name = os.path.splitext(os.path.basename(file))[0]

        # Load module dynamically
        spec = importlib.util.spec_from_file_location(module_name, file)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load spec for {file}")

        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)

        # Find FeatureStore instance in the module
        store = None
        for attr_name in dir(module):
            attr = getattr(module, attr_name)
            if isinstance(attr, FeatureStore):
                store = attr
                break

        if not store:
            console.print(
                "[bold red]Error:[/bold red] No FeatureStore instance found in file."
            )
            raise typer.Exit(code=1)

        # === Startup Checks (UX) ===
        # Check 1: RAG Usage without Keys
        has_retrievers = len(store.retriever_registry.retrievers) > 0
        has_openai = os.getenv("OPENAI_API_KEY") is not None
        has_cohere = os.getenv("COHERE_API_KEY") is not None

        if has_retrievers and not (has_openai or has_cohere):
            console.print(
                Panel(
                    "[yellow]Warning: Retrievers detected but no LLM API Key found.[/yellow]\n"
                    "Vector search and generation will fail.\n"
                    "Fix: set [bold]OPENAI_API_KEY[/bold] or [bold]COHERE_API_KEY[/bold] in .env",
                    title="Configuration Warning",
                    border_style="yellow",
                )
            )

        # Start scheduler
        store.start()

        # Set API key in env for the server to pick up
        if api_key:
            os.environ["MERIDIAN_API_KEY"] = api_key

        app = create_app(store)

        # Create Rich Layout
        layout = Layout()
        layout.split(
            Layout(name="header", size=3),
            Layout(name="main", ratio=1),
            Layout(name="footer", size=3),
        )

        layout["header"].update(
            Panel(
                f"Meridian Feature Store | Serving {file} on http://{host}:{port}",
                style="bold white on blue",
            )
        )

        # Create a table for metrics
        def generate_metrics_table() -> Panel:
            table = Table(title="Live Metrics")
            table.add_column("Metric", style="cyan")
            table.add_column("Value", style="magenta")

            # Display basic static metrics (Real-time dashboard requires Epic 4.2)
            # In a real implementation, we'd read from a shared metrics buffer
            table.add_row("Status", "Running 🟢")
            table.add_row("Uptime", datetime.now().strftime("%H:%M:%S"))
            table.add_row("Entities", str(len(store.registry.entities)))
            table.add_row("Features", str(len(store.registry.features)))
            return Panel(table, title="System Status", border_style="green")

        layout["main"].update(generate_metrics_table())

        layout["footer"].update(Panel("Press Ctrl+C to stop", style="dim"))

        console.print(f"[green]Successfully loaded features from {file}[/green]")
        console.print(
            "Starting server... (TUI disabled for simple log output compatibility)"
        )

        # NOTE: Running uvicorn inside a Rich Live context is tricky because uvicorn takes over stdout.
        # For this MVP, we will just print the rich header and then run uvicorn.
        # A full TUI requires running uvicorn in a separate thread/process and capturing logs.

        console.print(layout)
        uvicorn.run(app, host=host, port=port)

    except Exception as e:
        console.print(f"[bold red]Error loading features:[/bold red] {e}")
        raise typer.Exit(code=1)


@app.command(name="ui")
def ui(
    file: str = typer.Argument(
        ..., help="Path to the feature definition file (e.g., features.py)"
    ),
) -> None:
    """
    Launches the Meridian UI (Streamlit).
    """
    if not os.path.exists(file):
        console.print(f"[bold red]Error:[/bold red] File '{file}' not found.")
        raise typer.Exit(code=1)

    try:
        import streamlit.web.cli as stcli
    except ImportError:
        console.print(
            "[bold red]Error:[/bold red] Streamlit is not installed. "
            "Please run [green]pip install meridian[ui][/green] or [green]pip install streamlit[/green]."
        )
        raise typer.Exit(code=1)

    # Resolve absolute path to the UI script
    ui_script = os.path.join(os.path.dirname(__file__), "ui.py")

    # Construct the command args for streamlit
    # Format: streamlit run path/to/ui.py --server.headless=true --server.port=8501 -- path/to/features.py
    sys.argv = [
        "streamlit",
        "run",
        ui_script,
        "--server.headless=true",
        "--server.port=8501",
        "--",
        file,
    ]

    console.print(f"[green]Launching Meridian UI for {file}...[/green]")
    sys.exit(stcli.main())


@app.command(name="doctor")
def doctor() -> None:
    """
    Diagnose configuration and connectivity issues.
    """
    from .doctor import run_doctor

    run_doctor()


context_app = typer.Typer(help="Manage and inspect Context Store assemblies.")
app.add_typer(context_app, name="context")


@context_app.command(name="explain")
def explain_cmd(
    ctx_id: str = typer.Argument(..., help="The Context ID to trace"),
    host: str = typer.Option("127.0.0.1", help="Meridian server host"),
    port: int = typer.Option(8000, help="Meridian server port"),
) -> None:
    """
    Fetch and display a context trace (RAG explanation).
    """
    import urllib.request
    import urllib.error
    import json

    url = f"http://{host}:{port}/context/{ctx_id}/explain"
    console.print(f"Fetching trace for [bold cyan]{ctx_id}[/bold cyan] from {url}...")

    if not url.lower().startswith(("http://", "https://")):
        console.print("[bold red]Error:[/bold red] Invalid URL scheme")
        raise typer.Exit(1)

    try:
        # Use standard lib to avoid extra dependencies for CLI
        req = urllib.request.Request(url)
        # Add API Key if present in env
        api_key = os.getenv("MERIDIAN_API_KEY")
        if api_key:
            req.add_header("X-API-Key", api_key)

        with urllib.request.urlopen(req) as response:  # nosec B310
            if response.status != 200:
                console.print(
                    f"[bold red]Error:[/bold red] Server returned {response.status}"
                )
                raise typer.Exit(1)

            data = json.loads(response.read().decode())

            # Pretty print with Rich
            console.print(
                Panel(
                    json.dumps(data, indent=2),
                    title=f"Context Trace: {ctx_id}",
                    border_style="green",
                )
            )

    except urllib.error.URLError as e:
        console.print(
            f"[bold red]Connection Failed:[/bold red] {e}. Run [bold]meridian doctor[/bold] to check connectivity."
        )
        raise typer.Exit(1)
    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] {e}")
        raise typer.Exit(1)


@app.command(name="index")
def index_cmd(
    action: str = typer.Argument(..., help="Action: create, status"),
    name: str = typer.Argument(..., help="Name of the index"),
    dimension: int = typer.Option(1536, help="Vector dimension (create only)"),
    postgres_url: str = typer.Option(
        None, envvar="MERIDIAN_POSTGRES_URL", help="Postgres URL Override"
    ),
) -> None:
    """
    Manage vector indexes.
    Usage:
      meridian index create my_index --dimension=1536
      meridian index status my_index
    """
    import asyncio
    from .store.postgres import PostgresOfflineStore
    from sqlalchemy import text

    url = postgres_url or os.getenv("MERIDIAN_POSTGRES_URL")
    if not url:
        console.print("[bold red]Error:[/bold red] MERIDIAN_POSTGRES_URL not set.")
        raise typer.Exit(1)

    async def run_action() -> None:
        try:
            store = PostgresOfflineStore(url)

            if action == "create":
                console.print(
                    f"Creating index [bold]{name}[/bold] (dim={dimension})..."
                )
                await store.create_index_table(name, dimension)
                console.print(f"[green]Index '{name}' created successfully.[/green]")

            elif action == "status":
                table = f"meridian_index_{name}"
                async with store.engine.connect() as conn:  # type: ignore[no-untyped-call]
                    # Check if exists
                    exists = await conn.execute(
                        text(f"SELECT to_regclass('public.{table}')")
                    )
                    if not exists.scalar():
                        console.print(
                            f"[red]Index table '{table}' does not exist.[/red]"
                        )
                        return

                    # Count rows
                    res = await conn.execute(text(f"SELECT COUNT(*) FROM {table}"))  # nosec
                    count = res.scalar()
                    console.print(f"Index: [bold]{name}[/bold]")
                    console.print(f"Table: {table}")
                    console.print(f"Rows:  [cyan]{count}[/cyan]")

            else:
                console.print(f"[red]Unknown action: {action}[/red]")

            await store.engine.dispose()  # type: ignore[no-untyped-call]

        except Exception as e:
            console.print(f"[bold red]Error:[/bold red] {e}")
            raise typer.Exit(1)

    asyncio.run(run_action())


if __name__ == "__main__":
    app()
