"""Async implementation of nikame up with Rich CLI streaming."""

import asyncio
import json
import os
import sys
import time
from pathlib import Path
from typing import Any, Optional

from rich.console import Console
from rich.layout import Layout
from rich.live import Live
from rich.panel import Panel
from rich.table import Table
from rich.text import Text

console = Console()

class UpDisplay:
    def __init__(self, total_phases: int, project_name: str) -> None:
        self.total_phases = total_phases
        self.project_name = project_name
        self.current_phase = 0
        self.phase_name = "Initializing..."
        self.phase_start_time = time.time()
        
        self.logs: list[str] = []
        self.services: dict[str, dict] = {}
        self.finished = False

        self.layout = Layout()
        self.layout.split_column(
            Layout(name="header", size=3),
            Layout(name="services", ratio=1),
            Layout(name="logs", size=14)
        )

    def set_phase(self, phase_num: int, name: str) -> None:
        self.current_phase = phase_num
        self.phase_name = name
        self.phase_start_time = time.time()
        
    def add_log(self, text: str) -> None:
        # clean text
        text = text.strip()
        if not text:
            return
        self.logs.append(text)
        if len(self.logs) > 12:
            self.logs = self.logs[-12:]
            
    def update_service(self, name: str, status: str, health: str, endpoint: str = "internal") -> None:
        if name not in self.services:
            self.services[name] = {"start_time": time.time(), "endpoint": endpoint}
        self.services[name]["status"] = status
        self.services[name]["health"] = health
        # Update endpoint only if provided a real one
        if endpoint != "internal":
            self.services[name]["endpoint"] = endpoint

    def mark_finished(self) -> None:
        self.finished = True

    def render_header(self) -> Panel:
        elapsed = int(time.time() - self.phase_start_time)
        return Panel(
            f"Phase {self.current_phase} / {self.total_phases} — [bold]{self.phase_name}[/bold] (Elapsed: {elapsed}s)",
            style="cyan"
        )
        
    def render_services(self) -> Panel:
        if not self.services and not self.finished:
            return Panel(Text("Waiting for services...", style="dim italic"), title="[bold]Services[/bold]", border_style="blue")
            
        table = Table(expand=True)
        table.add_column("Service", style="cyan")
        table.add_column("Status")
        table.add_column("Health")
        table.add_column("Elapsed")
        table.add_column("Endpoint", style="green")
        
        for name, data in sorted(self.services.items()):
            st = data.get("status", "waiting")
            h = data.get("health", "waiting")
            el = int(time.time() - data.get("start_time", time.time()))
            
            st_text = "[yellow]🟡 starting[/yellow]" if st in ("starting", "created") else ("[green]🟢 running[/green]" if st == "running" else ("[red]🔴 failed[/red]" if st == "exited" else "[dim]⚪ waiting[/dim]"))
            h_text = "[yellow]🟡 starting[/yellow]" if h == "starting" else ("[green]🟢 healthy[/green]" if h in ("healthy", "N/A") else ("[red]🔴 failed[/red]" if h == "unhealthy" else f"[dim]⚪ {h}[/dim]"))
            
            table.add_row(name, st_text, h_text, f"{el}s", data.get("endpoint", "internal"))
            
        return Panel(table, title="[bold]Services[/bold]", border_style="blue")
        
    def render_logs(self) -> Panel:
        text = "\n".join(self.logs)
        return Panel(Text(text, style="dim"), title="[bold]Logs[/bold]", border_style="green")
        
    def render_quick_start(self) -> Panel:
        lines = [
            "[bold green]✨ All systems online and healthy! ✨[/bold green]",
            "",
            "[bold cyan]Next Steps:[/bold cyan]",
            "  1. Try the API: [bold]curl http://localhost:8000/docs[/bold]",
            "  2. Run tests: [bold]nikame test[/bold]",
            "  3. Check health: [bold]nikame doctor[/bold]",
            "",
            "[dim]Press [bold]Ctrl+C[/bold] to stop services or [bold]nikame down[/bold] later.[/dim]"
        ]
        return Panel("\n".join(lines), title="[bold]🚀 Quick Start[/bold]", border_style="bright_magenta", padding=(1, 2))

    def __rich__(self):
        if self.finished:
            return self.render_quick_start()
            
        self.layout["header"].update(self.render_header())
        self.layout["services"].update(self.render_services())
        self.layout["logs"].update(self.render_logs())
        return self.layout

async def run_logged_cmd(cmd: list[str], display: UpDisplay, cwd: Path, timeout: int, throw_on_error: bool = True):
    """Run a command, stream its stdout/stderr to display logs, and enforce total timeout."""
    import subprocess
    process = await asyncio.create_subprocess_exec(
        *cmd,
        cwd=str(cwd),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        # ensure environment variables are passed
        env=os.environ.copy()
    )
    
    async def _read_stream(stream) -> None:
        while True:
            line = await stream.readline()
            if not line:
                break
            # Decode and clean the line
            decoded = line.decode('utf-8', errors='replace').rstrip()
            display.add_log(decoded)
            
    # Read output
    try:
        await asyncio.wait_for(
            asyncio.gather(_read_stream(process.stdout)),
            timeout=timeout
        )
        await process.wait()
    except TimeoutError:
        try:
            process.terminate()
        except Exception:
            pass
        if throw_on_error:
            raise RuntimeError(f"Command timed out after {timeout}s: {' '.join(cmd)}")
    
    if process.returncode != 0 and throw_on_error:
        raise RuntimeError(f"Command failed with code {process.returncode}: {' '.join(cmd)}")
    return process.returncode

async def poll_health(service_name: str, display: UpDisplay, timeout: int) -> bool:
    """Poll docker inspect for a service until it becomes healthy or times out."""
    start = time.time()
    # Find the actual container name (could be project-service-1)
    container_name = f"{display.project_name}-{service_name}-1"
    
    while time.time() - start < timeout:
        # Check status and health
        cmd = ["docker", "inspect", "--format='{{.State.Status}}|{{.State.Health.Status}}'", container_name]
        proc = await asyncio.create_subprocess_exec(
            *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
        )
        stdout, _ = await proc.communicate()
        output = stdout.decode().strip().strip("'")
        
        if proc.returncode == 0 and output:
            parts = output.split("|")
            status = parts[0]
            health = parts[1] if len(parts) > 1 and parts[1] != "<no value>" else "N/A"
            
            display.update_service(service_name, status, health)
            
            if status == "exited":
                return False
            if health == "healthy" or (health == "N/A" and status == "running"):
                # N/A health means it doesn't have a healthcheck, but it's running
                return True
        else:
            display.update_service(service_name, "missing", "waiting")
            
        await asyncio.sleep(2)
        
    return False

async def execute_up_async(
    project_dir: Path, 
    config: Any, 
    service: tuple[str, ...], 
    build: bool, 
    detach: bool, 
    compose_file: Path, 
    compose_data: dict, 
    cmd_base: list[str],
    timeout: int = 300
) -> None:
    # Build phase
    total_phases = 5
    display = UpDisplay(total_phases=total_phases, project_name=config.name)
    
    with Live(display, refresh_per_second=4, console=console) as live:
        
        # Phase 1: Build
        if build:
            display.set_phase(1, "Building Images")
            cmd = ["docker", "buildx", "bake", "-f", str(compose_file), "--progress=plain"]
            try:
                await run_logged_cmd(cmd, display, project_dir, timeout=300)
            except Exception as e:
                console.print(f"[bold red]Build failed:[/bold red] {e}")
                raise SystemExit(1)
        
        # Outline phases
        all_services = list(compose_data.get("services", {}).keys())
        # Populate initial display
        for s in all_services:
            display.update_service(s, "waiting", "waiting")
            
        phase1_keywords = {"postgres", "pgbouncer", "dragonfly", "redis", "minio", "qdrant", "weaviate", "milvus", "chroma", "redpanda", "kafka", "clickhouse", "mongodb", "neo4j"}
        phase2_keywords = {"mlflow", "langfuse", "evidently", "prefect", "airflow", "grafana", "prometheus", "alertmanager"}
        phase3_keywords = {"vllm", "ollama", "llamacpp", "tgi", "triton", "localai", "xinference", "airllm", "bentoml", "whisper", "tts"}
        
        phase1 = [s for s in all_services if any(k in s for k in phase1_keywords)]
        phase2 = [s for s in all_services if any(k in s for k in phase2_keywords) and s not in phase1]
        phase3 = [s for s in all_services if any(k in s for k in phase3_keywords) and s not in phase1 and s not in phase2]
        remaining = [s for s in all_services if s not in phase1 and s not in phase2 and s not in phase3]
        
        # Scale default timeouts based on user input
        scale = timeout / 300.0 if timeout != 300 else 1.0
        
        phases = [
            ("Starting Data Stores", phase1, int(1200 * scale)),
            ("Starting Observability", phase2, int(300 * scale)),
            ("Starting Application API", remaining, int(600 * scale)),
            ("Starting LLM Engines", phase3, int(1800 * scale)),
        ]
        
        phase_idx = 1 if not build else 2
        for label, svcs, phase_timeout in phases:
            if not svcs: continue
            
            # Start Ollama pull in background if Ollama is in this phase
            if "ollama" in svcs:
                # Fire and forget ollama pull. The API should handle 503s.
                display.add_log("Starting background Ollama model pull (non-blocking)...")
                # We start the container first then run exec, but wait, 'up -d' starts it
                pass # Handled after container starts
                
            display.set_phase(phase_idx, label)
            phase_cmd = [*cmd_base, "up", "-d", *svcs]
            try:
                await run_logged_cmd(phase_cmd, display, project_dir, timeout=phase_timeout)
                
                if "ollama" in svcs:
                    # Run background pull
                    model = config.ml_serving.get("ollama", {}).get("model", "nomic-embed-text")
                    # Use Popen to completely detach, or asyncio.create_subprocess_exec without awaiting
                    # But we wait just a moment for ollama to be up
                    import subprocess
                    subprocess.Popen(["docker", "exec", f"{config.name}-ollama-1", "ollama", "pull", model], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                    display.add_log(f"Triggered background pull for model: {model}")

            except Exception as e:
                console.print(f"[bold red]Phase Failed ({label}):[/bold red] {e}")
                raise SystemExit(1)
            phase_idx += 1
            
        display.set_phase(total_phases, "Health Verification")
        # Concurrent health poll for all services
        tasks = [poll_health(svc, display, timeout=60) for svc in all_services]
        results = await asyncio.gather(*tasks)
        
        if not all(results):
            display.add_log("ERROR: Some services failed to become healthy.")
            console.print("\n[bold red]Health check failed for some services.[/bold red]")
            # Identify which ones
            failed = [s for s, r in zip(all_services, results, strict=False) if not r]
            console.print(f"Failed services: {', '.join(failed)}")
            raise SystemExit(1)
            
        # Verify postgres/pgbouncer explicit health before migrations (covered by all(results))
        # Now run migrations!
        app_dir = project_dir / "app"
        if (app_dir / "alembic").exists() and (not service or "api" in service):
            display.set_phase(total_phases, "Database Migrations")
            cmd = ["docker", "compose", "-p", config.name, "-f", str(compose_file), "run", "--rm", "api", "alembic", "upgrade", "head"]
            try:
                await run_logged_cmd(cmd, display, project_dir, timeout=60)
            except Exception as e:
                console.print(f"[bold red]Migration failed:[/bold red] {e}")
                raise SystemExit(1)
                
            # Seed DB conditionally
            seed_script = project_dir / "scripts" / "seed_db.py"
            if seed_script.exists():
                display.add_log("Seeding database...")
                try:
                    seed_cmd = ["docker", "compose", "-p", config.name, "-f", str(compose_file), "run", "--rm", "api", "python", "/app/../scripts/seed_db.py"]
                    await run_logged_cmd(seed_cmd, display, project_dir, timeout=60)
                except Exception:
                    display.add_log("Seeding failed (non-fatal)")
        
        display.mark_finished()
        display.set_phase(total_phases, "All Services Online")
        await asyncio.sleep(1) # Final render update
