import re
import secrets
import shlex
import sqlite3
import subprocess
import uuid
from datetime import datetime, timezone
from pathlib import Path


DB_PATH = Path.home() / ".onako" / "onako.db"


class TmuxOrchestrator:
    def __init__(self, session_name: str = "onako", db_path: Path | None = None):
        self.session_name = session_name
        self.db_path = db_path or DB_PATH
        self.tasks: dict[str, dict] = {}
        self._window_ids: dict[str, str] = {}  # window_name -> @id
        self._last_content_length: dict[str, int] = {}  # window_name -> pane content length
        self._last_output_cache: dict[str, str] = {}  # window_name -> cached stripped output
        self._init_db()
        self._load_tasks()
        self._ensure_session()
        self.rediscover_tasks()

    @property
    def _session_target(self) -> str:
        """Session name with colon suffix for unambiguous tmux targeting."""
        return f"{self.session_name}:"

    def _ensure_session(self):
        result = subprocess.run(
            ["tmux", "has-session", "-t", self._session_target],
            capture_output=True,
        )
        if result.returncode != 0:
            subprocess.run(
                ["tmux", "new-session", "-d", "-s", self.session_name, "-n", "onako-main"],
                check=True,
            )

    def _init_db(self):
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(self.db_path)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS tasks (
                id TEXT PRIMARY KEY,
                prompt TEXT,
                status TEXT,
                started_at TEXT,
                origin TEXT DEFAULT 'managed',
                last_activity TEXT,
                claude_session_id TEXT,
                branch TEXT,
                worktree_path TEXT
            )
        """)
        # Migrate existing DBs that lack newer columns
        cursor = conn.execute("PRAGMA table_info(tasks)")
        columns = [row[1] for row in cursor.fetchall()]
        if "origin" not in columns:
            conn.execute("ALTER TABLE tasks ADD COLUMN origin TEXT DEFAULT 'managed'")
        if "last_activity" not in columns:
            conn.execute("ALTER TABLE tasks ADD COLUMN last_activity TEXT")
        if "claude_session_id" not in columns:
            conn.execute("ALTER TABLE tasks ADD COLUMN claude_session_id TEXT")
        if "branch" not in columns:
            conn.execute("ALTER TABLE tasks ADD COLUMN branch TEXT")
        if "worktree_path" not in columns:
            conn.execute("ALTER TABLE tasks ADD COLUMN worktree_path TEXT")
        if "last_output" not in columns:
            conn.execute("ALTER TABLE tasks ADD COLUMN last_output TEXT")
        conn.commit()

        # Events table for stats tracking
        conn.execute("""
            CREATE TABLE IF NOT EXISTS events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                event_type TEXT NOT NULL,
                task_id TEXT,
                timestamp TEXT NOT NULL,
                metadata TEXT
            )
        """)

        # Backfill events from existing tasks (only if events table is empty)
        event_count = conn.execute("SELECT COUNT(*) FROM events").fetchone()[0]
        if event_count == 0:
            tasks = conn.execute("SELECT id, status, started_at, last_activity FROM tasks").fetchall()
            for task_id, status, started_at, last_activity in tasks:
                # Every task gets a task_created event
                conn.execute(
                    "INSERT INTO events (event_type, task_id, timestamp) VALUES (?, ?, ?)",
                    ("task_created", task_id, started_at or datetime.now(timezone.utc).isoformat()),
                )
                # Completed tasks get a task_completed event
                if status in ("done", "hidden"):
                    conn.execute(
                        "INSERT INTO events (event_type, task_id, timestamp) VALUES (?, ?, ?)",
                        ("task_completed", task_id, last_activity or started_at or datetime.now(timezone.utc).isoformat()),
                    )

        conn.commit()
        conn.close()

    def _load_tasks(self):
        conn = sqlite3.connect(self.db_path)
        rows = conn.execute("SELECT id, prompt, status, started_at, origin, last_activity, claude_session_id, branch, worktree_path, last_output FROM tasks").fetchall()
        conn.close()
        for row in rows:
            self.tasks[row[0]] = {
                "id": row[0],
                "prompt": row[1],
                "status": row[2],
                "started_at": row[3],
                "origin": row[4] or "managed",
                "last_activity": row[5],
                "claude_session_id": row[6],
                "branch": row[7],
                "worktree_path": row[8],
                "last_output": row[9],
            }

    def _save_task(self, task: dict):
        conn = sqlite3.connect(self.db_path)
        conn.execute(
            "INSERT OR REPLACE INTO tasks (id, prompt, status, started_at, origin, last_activity, claude_session_id, branch, worktree_path, last_output) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (task["id"], task["prompt"], task["status"], task["started_at"], task.get("origin", "managed"), task.get("last_activity"), task.get("claude_session_id"), task.get("branch"), task.get("worktree_path"), task.get("last_output")),
        )
        conn.commit()
        conn.close()

    def _log_event(self, event_type: str, task_id: str | None = None, metadata: str | None = None):
        conn = sqlite3.connect(self.db_path)
        conn.execute(
            "INSERT INTO events (event_type, task_id, timestamp, metadata) VALUES (?, ?, ?, ?)",
            (event_type, task_id, datetime.now(timezone.utc).isoformat(), metadata),
        )
        conn.commit()
        conn.close()

    def _run_tmux(self, *args) -> subprocess.CompletedProcess:
        return subprocess.run(["tmux", *args], capture_output=True, text=True)

    def _task_target(self, task_id: str) -> str:
        """Return a safe tmux target for a task, using window ID when available."""
        window_id = self._window_ids.get(task_id)
        if window_id:
            return window_id
        return f"{self.session_name}:{task_id}"

    def create_task(self, command: str, working_dir: str | None = None, prompt: str | None = None, claude_session_id: str | None = None, branch: str | None = None, worktree_path: str | None = None) -> dict:
        self._ensure_session()
        task_id = f"task-{secrets.token_hex(4)}"
        self._run_tmux(
            "new-window", "-t", self._session_target, "-n", task_id,
        )
        # Look up the window ID for safe targeting
        result = self._run_tmux(
            "list-windows", "-t", self._session_target, "-F", "#{window_name}|#{window_id}",
        )
        if result.stdout.strip():
            for line in result.stdout.strip().split("\n"):
                name, _, wid = line.partition("|")
                if name == task_id and wid:
                    self._window_ids[task_id] = wid
                    break
        if working_dir:
            self._run_tmux(
                "send-keys", "-t", self._task_target(task_id),
                f"cd {shlex.quote(working_dir)}", "Enter",
            )
        self._run_tmux(
            "send-keys", "-t", self._task_target(task_id),
            "-l", command,
        )
        self._run_tmux(
            "send-keys", "-t", self._task_target(task_id),
            "Enter",
        )
        now = datetime.now(timezone.utc).isoformat()
        task = {
            "id": task_id,
            "prompt": prompt or command,
            "status": "running",
            "started_at": now,
            "origin": "managed",
            "last_activity": now,
            "claude_session_id": claude_session_id,
            "branch": branch,
            "worktree_path": worktree_path,
        }
        self.tasks[task_id] = task
        self._save_task(task)
        self._log_event("task_created", task_id)
        return task

    def resume_task(self, task_id: str, skip_permissions: bool = False, worktree_manager=None) -> dict:
        """Resume a done task by reopening its tmux window."""
        task = self.tasks.get(task_id)
        if not task:
            raise ValueError("Task not found")
        if task["status"] == "running":
            raise ValueError("Task is still running")
        session_id = task.get("claude_session_id")
        if not session_id:
            raise ValueError("Task has no claude_session_id (cannot resume)")
        # Re-create worktree if it was cleaned up
        if task.get("worktree_path") and not Path(task["worktree_path"]).exists():
            if worktree_manager and task.get("branch"):
                worktree_manager.ensure_worktree(task["branch"])
        cmd = "claude"
        if skip_permissions:
            cmd += " --dangerously-skip-permissions"
        cmd += f" --resume {session_id}"
        # Create a new tmux window with the original task name
        self._ensure_session()
        self._run_tmux("new-window", "-t", self._session_target, "-n", task_id)
        # Look up the window ID for safe targeting
        result = self._run_tmux(
            "list-windows", "-t", self._session_target, "-F", "#{window_name}|#{window_id}",
        )
        if result.stdout.strip():
            for line in result.stdout.strip().split("\n"):
                name, _, wid = line.partition("|")
                if name == task_id and wid:
                    self._window_ids[task_id] = wid
                    break
        # cd to worktree if task was running in one
        if task.get("worktree_path"):
            self._run_tmux(
                "send-keys", "-t", self._task_target(task_id),
                f"cd {shlex.quote(task['worktree_path'])}", "Enter",
            )
        self._run_tmux("send-keys", "-t", self._task_target(task_id), "-l", cmd)
        self._run_tmux("send-keys", "-t", self._task_target(task_id), "Enter")
        now = datetime.now(timezone.utc).isoformat()
        task["status"] = "running"
        task["last_output"] = None
        task["last_activity"] = now
        if not task["prompt"].startswith("(resumed)"):
            task["prompt"] = f"(resumed) {task['prompt']}"
        self._save_task(task)
        self._log_event("task_resumed", task_id)
        return task

    def list_tasks(self) -> list[dict]:
        self.rediscover_tasks()
        self._sync_task_status()
        return list(self.tasks.values())

    _ANSI_RE = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]")

    def get_output(self, task_id: str) -> str:
        raw = self.get_raw_output(task_id)
        return self._strip_claude_chrome(raw)

    @classmethod
    def _strip_claude_chrome(cls, text: str) -> str:
        lines = text.split("\n")
        # Strip from the bottom: Claude Code's TUI draws box-drawing chars,
        # the › prompt, and status lines like "accept edits on..."
        while lines:
            line = cls._ANSI_RE.sub("", lines[-1]).strip()
            if (
                not line
                or all(c in "─━╭╮╰╯│┃┌┐└┘├┤┬┴┼╋═║ ›❯▸▶" for c in line)
                or "accept edits" in line
                or "esc to interrupt" in line
                or "shift+tab to cycle" in line
                or "ctrl+" in line.lower()
                or line == "›"
            ):
                lines.pop()
            else:
                break
        return "\n".join(lines)

    def get_raw_output(self, task_id: str) -> str:
        result = self._run_tmux(
            "capture-pane", "-t", self._task_target(task_id),
            "-p", "-S", "-", "-e",
        )
        return result.stdout

    def send_message(self, task_id: str, message: str):
        self._run_tmux(
            "send-keys", "-t", self._task_target(task_id),
            "-l", message,
        )
        self._run_tmux(
            "send-keys", "-t", self._task_target(task_id),
            "Enter",
        )
        self._log_event("message_sent", task_id)

    def send_interrupt(self, task_id: str):
        self._run_tmux("send-keys", "-t", self._task_target(task_id), "Escape")
        self._log_event("task_interrupted", task_id)

    def window_count(self) -> int:
        result = self._run_tmux(
            "list-windows", "-t", self._session_target, "-F", "#{window_name}",
        )
        if result.stdout.strip():
            return len(result.stdout.strip().split("\n"))
        return 0

    def kill_task(self, task_id: str):
        # Snapshot output before killing the window
        try:
            output = self.get_output(task_id)
            lines = output.split("\n")
            snapshot = "\n".join(lines[-200:]) if len(lines) > 200 else output
        except Exception:
            snapshot = None
        self._run_tmux("kill-window", "-t", self._task_target(task_id))
        if task_id in self.tasks:
            self.tasks[task_id]["status"] = "done"
            if snapshot:
                self.tasks[task_id]["last_output"] = snapshot
            self._save_task(self.tasks[task_id])
            self._log_event("task_killed", task_id)

    def hide_task(self, task_id: str):
        """Hide a done task from the dashboard."""
        task = self.tasks.get(task_id)
        if not task:
            raise ValueError("Task not found")
        if task["status"] != "done":
            raise ValueError("Task is not done")
        task["status"] = "hidden"
        self._save_task(task)

    def _sync_task_status(self):
        result = self._run_tmux(
            "list-windows", "-t", self._session_target, "-F", "#{window_name}|#{window_id}",
        )
        active_windows = set()
        if result.stdout.strip():
            for line in result.stdout.strip().split("\n"):
                parts = line.split("|", 1)
                active_windows.add(parts[0])
                if len(parts) > 1:
                    self._window_ids[parts[0]] = parts[1]
        for task_id, task in self.tasks.items():
            if task["status"] == "running" and task_id not in active_windows:
                task["status"] = "done"
                cached = self._last_output_cache.pop(task_id, None)
                if cached:
                    task["last_output"] = cached
                self._save_task(task)
                self._log_event("task_completed", task_id)
            elif task["status"] == "running" and task_id in active_windows:
                pane = self._run_tmux(
                    "capture-pane", "-t", self._task_target(task_id), "-p", "-S", "-", "-e",
                )
                raw = pane.stdout
                content_len = len(raw)
                prev_len = self._last_content_length.get(task_id)
                self._last_content_length[task_id] = content_len
                # Cache stripped output for snapshot on natural finish
                stripped = self._strip_claude_chrome(raw)
                lines = stripped.split("\n")
                self._last_output_cache[task_id] = "\n".join(lines[-200:]) if len(lines) > 200 else stripped
                if prev_len is not None and content_len != prev_len:
                    task["last_activity"] = datetime.now(timezone.utc).isoformat()
                    self._save_task(task)

    def get_stats(self) -> dict:
        """Return aggregated stats from the events table."""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row

        # --- Volume ---
        counts = conn.execute("""
            SELECT
                COALESCE(SUM(CASE WHEN event_type = 'task_created' THEN 1 ELSE 0 END), 0) AS total_tasks,
                COALESCE(SUM(CASE WHEN event_type = 'message_sent' THEN 1 ELSE 0 END), 0) AS total_messages_sent,
                COALESCE(SUM(CASE WHEN event_type = 'task_resumed' THEN 1 ELSE 0 END), 0) AS total_resumes,
                COALESCE(SUM(CASE WHEN event_type = 'task_killed' THEN 1 ELSE 0 END), 0) AS total_kills,
                COALESCE(SUM(CASE WHEN event_type = 'task_interrupted' THEN 1 ELSE 0 END), 0) AS total_interrupts
            FROM events
        """).fetchone()

        # Runtime: sum durations from task_created to task_completed (or now) per task
        created_rows = conn.execute(
            "SELECT task_id, timestamp FROM events WHERE event_type = 'task_created'"
        ).fetchall()
        completed_rows = conn.execute(
            "SELECT task_id, timestamp FROM events WHERE event_type = 'task_completed'"
        ).fetchall()
        completed_map: dict[str, str] = {}
        for row in completed_rows:
            completed_map[row["task_id"]] = row["timestamp"]

        now = datetime.now(timezone.utc)
        total_runtime = 0
        for row in created_rows:
            task_id = row["task_id"]
            start = datetime.fromisoformat(row["timestamp"])
            if task_id in completed_map:
                end = datetime.fromisoformat(completed_map[task_id])
            else:
                end = now
            total_runtime += int((end - start).total_seconds())

        volume = {
            "total_tasks": counts["total_tasks"],
            "total_messages_sent": counts["total_messages_sent"],
            "total_resumes": counts["total_resumes"],
            "total_kills": counts["total_kills"],
            "total_interrupts": counts["total_interrupts"],
            "total_runtime_seconds": total_runtime,
        }

        # --- Records ---
        # Longest task: task with max duration between created and completed
        longest_task = None
        if created_rows and completed_map:
            best_id = None
            best_duration = -1
            for row in created_rows:
                tid = row["task_id"]
                if tid in completed_map:
                    start = datetime.fromisoformat(row["timestamp"])
                    end = datetime.fromisoformat(completed_map[tid])
                    dur = int((end - start).total_seconds())
                    if dur > best_duration:
                        best_duration = dur
                        best_id = tid
            if best_id is not None:
                prompt = self.tasks.get(best_id, {}).get("prompt", "")
                longest_task = {"id": best_id, "prompt": prompt, "duration_seconds": best_duration}

        # Most resumed task
        most_resumed_task = None
        resumed_row = conn.execute(
            "SELECT task_id, COUNT(*) as cnt FROM events WHERE event_type = 'task_resumed' GROUP BY task_id ORDER BY cnt DESC LIMIT 1"
        ).fetchone()
        if resumed_row:
            tid = resumed_row["task_id"]
            prompt = self.tasks.get(tid, {}).get("prompt", "")
            most_resumed_task = {"id": tid, "prompt": prompt, "resume_count": resumed_row["cnt"]}

        # Busiest day (most task_created events)
        busiest_day = None
        busiest_row = conn.execute(
            "SELECT SUBSTR(timestamp, 1, 10) AS date, COUNT(*) AS cnt FROM events WHERE event_type = 'task_created' GROUP BY date ORDER BY cnt DESC LIMIT 1"
        ).fetchone()
        if busiest_row:
            busiest_day = {"date": busiest_row["date"], "task_count": busiest_row["cnt"]}

        # Biggest day for messages
        biggest_day_messages = None
        biggest_msg_row = conn.execute(
            "SELECT SUBSTR(timestamp, 1, 10) AS date, COUNT(*) AS cnt FROM events WHERE event_type = 'message_sent' GROUP BY date ORDER BY cnt DESC LIMIT 1"
        ).fetchone()
        if biggest_msg_row:
            biggest_day_messages = {"date": biggest_msg_row["date"], "message_count": biggest_msg_row["cnt"]}

        records = {
            "longest_task": longest_task,
            "most_resumed_task": most_resumed_task,
            "busiest_day": busiest_day,
            "biggest_day_messages": biggest_day_messages,
        }

        # --- Trends (ISO week grouping, done in Python for correctness) ---
        task_created_rows = conn.execute(
            "SELECT timestamp FROM events WHERE event_type = 'task_created'"
        ).fetchall()
        message_sent_rows = conn.execute(
            "SELECT timestamp FROM events WHERE event_type = 'message_sent'"
        ).fetchall()

        tasks_by_week: dict[str, int] = {}
        for row in task_created_rows:
            dt = datetime.fromisoformat(row["timestamp"])
            iso_year, iso_week, _ = dt.isocalendar()
            week_key = f"{iso_year}-W{iso_week:02d}"
            tasks_by_week[week_key] = tasks_by_week.get(week_key, 0) + 1

        messages_by_week: dict[str, int] = {}
        for row in message_sent_rows:
            dt = datetime.fromisoformat(row["timestamp"])
            iso_year, iso_week, _ = dt.isocalendar()
            week_key = f"{iso_year}-W{iso_week:02d}"
            messages_by_week[week_key] = messages_by_week.get(week_key, 0) + 1

        trends = {
            "tasks_per_week": [{"week": w, "count": c} for w, c in sorted(tasks_by_week.items())],
            "messages_per_week": [{"week": w, "count": c} for w, c in sorted(messages_by_week.items())],
        }

        # --- Heatmap: event count per calendar day ---
        heatmap_rows = conn.execute(
            "SELECT SUBSTR(timestamp, 1, 10) AS date, COUNT(*) AS cnt FROM events GROUP BY date ORDER BY date"
        ).fetchall()
        heatmap = {row["date"]: row["cnt"] for row in heatmap_rows}

        # --- Activity by hour: 24-element array ---
        activity_by_hour = [0] * 24
        hour_rows = conn.execute("SELECT timestamp FROM events").fetchall()
        for row in hour_rows:
            dt = datetime.fromisoformat(row["timestamp"])
            activity_by_hour[dt.hour] += 1

        conn.close()

        return {
            "volume": volume,
            "records": records,
            "trends": trends,
            "heatmap": heatmap,
            "activity_by_hour": activity_by_hour,
        }

    def rediscover_tasks(self):
        """Rediscover tasks from existing tmux windows on server restart."""
        result = self._run_tmux(
            "list-windows", "-t", self._session_target, "-F", "#{window_name}|#{window_id}",
        )
        if not result.stdout.strip():
            return
        for line in result.stdout.strip().split("\n"):
            parts = line.split("|", 1)
            window_name = parts[0]
            window_id = parts[1] if len(parts) > 1 else None
            if window_id:
                self._window_ids[window_name] = window_id
            if window_name not in self.tasks:
                is_managed = window_name.startswith("task-") or window_name == "onako-main"
                now = datetime.now(timezone.utc).isoformat()
                task = {
                    "id": window_name,
                    "prompt": "(rediscovered)" if is_managed else window_name,
                    "status": "running",
                    "started_at": None,
                    "origin": "managed" if is_managed else "external",
                    "last_activity": now,
                    "branch": None,
                    "worktree_path": None,
                }
                self.tasks[window_name] = task
                self._save_task(task)
            else:
                changed = False
                if self.tasks[window_name]["status"] == "done":
                    self.tasks[window_name]["status"] = "running"
                    changed = True
                if window_name == "onako-main" and self.tasks[window_name]["origin"] == "external":
                    self.tasks[window_name]["origin"] = "managed"
                    changed = True
                if not self.tasks[window_name].get("last_activity"):
                    self.tasks[window_name]["last_activity"] = datetime.now(timezone.utc).isoformat()
                    changed = True
                if changed:
                    self._save_task(self.tasks[window_name])
