"""
agentpaas.db.models — SQLAlchemy ORM models.

Core tables: tenants, agents, agent_versions, runs, jobs, api_keys, usage_records.
Principle P3: Works with both SQLite and PostgreSQL.
"""
from __future__ import annotations
import uuid
import hashlib
import json
import time
from datetime import datetime, timezone
from typing import Optional, Dict, Any, List

import sqlite3


def now_utc():
    return datetime.now(timezone.utc).isoformat()


def gen_id(prefix: str = "") -> str:
    return f"{prefix}{uuid.uuid4().hex[:12]}"


class Database:
    """
    Lightweight database abstraction. Uses SQLite directly for MVP.
    Designed for easy migration to SQLAlchemy/PostgreSQL later.
    """

    def __init__(self, db_url: str = "sqlite:///./agentpaas.db"):
        if "sqlite" in db_url:
            db_path = db_url.replace("sqlite:///", "")
            self.conn = sqlite3.connect(db_path, check_same_thread=False)
            self.conn.row_factory = sqlite3.Row
        else:
            raise NotImplementedError("Only SQLite supported in MVP. Set AGENTPAAS_DATABASE_URL=sqlite:///./agentpaas.db")
        self._init_tables()

    def _init_tables(self):
        c = self.conn.cursor()
        # Step 1: Create tables (IF NOT EXISTS — no-op for existing DBs)
        c.executescript("""
            CREATE TABLE IF NOT EXISTS tenants (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                plan TEXT DEFAULT 'free',
                quota_tokens_monthly INTEGER DEFAULT 1000000,
                quota_concurrency INTEGER DEFAULT 5,
                quota_agents INTEGER DEFAULT 20,
                status TEXT DEFAULT 'active',
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS users (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                email TEXT,
                role TEXT DEFAULT 'developer',
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS api_keys (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                user_id TEXT REFERENCES users(id),
                key_hash TEXT NOT NULL UNIQUE,
                key_prefix TEXT NOT NULL,
                name TEXT,
                scopes TEXT DEFAULT '[]',
                rate_limit INTEGER DEFAULT 60,
                expires_at TEXT,
                last_used_at TEXT,
                status TEXT DEFAULT 'active',
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS agents (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                name TEXT NOT NULL,
                description TEXT DEFAULT '',
                current_version INTEGER DEFAULT 1,
                tags TEXT DEFAULT '[]',
                environment TEXT DEFAULT 'production',
                traffic_rules TEXT,
                agent_dir TEXT,
                agent_template TEXT,
                instance_dir TEXT,
                status TEXT DEFAULT 'active',
                created_at TEXT,
                updated_at TEXT
            );

            CREATE TABLE IF NOT EXISTS agent_versions (
                agent_id TEXT REFERENCES agents(id),
                version INTEGER NOT NULL,
                config TEXT NOT NULL,
                config_hash TEXT NOT NULL,
                changelog TEXT DEFAULT '',
                created_by TEXT,
                created_at TEXT,
                PRIMARY KEY (agent_id, version)
            );

            CREATE TABLE IF NOT EXISTS runs (
                id TEXT PRIMARY KEY,
                agent_id TEXT REFERENCES agents(id),
                agent_version INTEGER,
                tenant_id TEXT REFERENCES tenants(id),
                job_id TEXT,
                input TEXT,
                output TEXT,
                status TEXT DEFAULT 'running',
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                steps INTEGER DEFAULT 0,
                duration_ms INTEGER DEFAULT 0,
                error TEXT,
                trace_id TEXT,
                workspace_path TEXT,
                idempotency_key TEXT,
                cancel_requested INTEGER DEFAULT 0,
                created_at TEXT,
                completed_at TEXT
            );

            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                agent_id TEXT REFERENCES agents(id),
                priority INTEGER DEFAULT 5,
                status TEXT DEFAULT 'queued',
                input TEXT,
                worker_id TEXT,
                attempts INTEGER DEFAULT 0,
                max_attempts INTEGER DEFAULT 3,
                started_at TEXT,
                completed_at TEXT,
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS usage_records (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                agent_id TEXT REFERENCES agents(id),
                run_id TEXT REFERENCES runs(id),
                model TEXT,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                tool_calls INTEGER DEFAULT 0,
                duration_ms INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0.0,
                created_at TEXT
            );

            -- 会话记忆 P1（docs/MEMORY_DESIGN.md §2.1）：一条 thread = 一段
            -- 对话；run 经 runs.thread_id 归属。summary_upto_at 是可排序游标
            -- （评审#2：用 run.created_at 而非 run_id）。
            CREATE TABLE IF NOT EXISTS threads (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                agent_id TEXT REFERENCES agents(id),
                title TEXT DEFAULT '',
                rolling_summary TEXT DEFAULT '',
                summary_upto_at TEXT DEFAULT '',
                status TEXT DEFAULT 'active',
                created_at TEXT, updated_at TEXT
            );

            -- 记忆操作成本单列（评审#3）：压缩/提炼成本不混进 runs。
            CREATE TABLE IF NOT EXISTS memory_ops (
                id TEXT PRIMARY KEY,
                tenant_id TEXT,
                thread_id TEXT,
                kind TEXT,                  -- compress | distill
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0.0,
                created_at TEXT
            );
        """)
        self.conn.commit()

        # Step 2: Migrations — add missing columns to existing tables
        self._migrate()

        # Step 3: Indexes — must run AFTER migrations so all columns exist
        c.executescript("""
            CREATE INDEX IF NOT EXISTS idx_agents_tenant ON agents(tenant_id);
            CREATE INDEX IF NOT EXISTS idx_agents_status ON agents(status);
            CREATE INDEX IF NOT EXISTS idx_agents_tenant_status ON agents(tenant_id, status);
            CREATE INDEX IF NOT EXISTS idx_runs_agent ON runs(agent_id);
            CREATE INDEX IF NOT EXISTS idx_runs_tenant ON runs(tenant_id);
            CREATE INDEX IF NOT EXISTS idx_runs_status ON runs(status);
            CREATE INDEX IF NOT EXISTS idx_runs_agent_tenant ON runs(agent_id, tenant_id);
            CREATE INDEX IF NOT EXISTS idx_runs_idempotency ON runs(idempotency_key);
            CREATE INDEX IF NOT EXISTS idx_jobs_tenant ON jobs(tenant_id);
            CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status);
            CREATE INDEX IF NOT EXISTS idx_jobs_agent ON jobs(agent_id);
            CREATE INDEX IF NOT EXISTS idx_usage_tenant ON usage_records(tenant_id);
            CREATE INDEX IF NOT EXISTS idx_usage_agent ON usage_records(agent_id);
            CREATE INDEX IF NOT EXISTS idx_usage_run ON usage_records(run_id);
            CREATE INDEX IF NOT EXISTS idx_api_keys_prefix ON api_keys(key_prefix);
            CREATE INDEX IF NOT EXISTS idx_threads_tenant_agent_updated
                ON threads(tenant_id, agent_id, updated_at DESC);
            CREATE INDEX IF NOT EXISTS idx_runs_thread ON runs(thread_id, created_at);
        """)
        self.conn.commit()

    def _ensure_table(self, name: str, ddl: str):
        try:
            self.conn.execute(f"SELECT 1 FROM {name} LIMIT 0")
        except sqlite3.OperationalError:
            self.conn.executescript(ddl)
            self.conn.commit()

    def _migrate(self):
        """Add columns to existing tables if missing (safe for fresh DBs too)."""
        migrations = [
            ("agents", "agent_dir", "ALTER TABLE agents ADD COLUMN agent_dir TEXT"),
            ("agents", "agent_template", "ALTER TABLE agents ADD COLUMN agent_template TEXT"),
            ("agents", "instance_dir", "ALTER TABLE agents ADD COLUMN instance_dir TEXT"),
            ("runs", "workspace_path", "ALTER TABLE runs ADD COLUMN workspace_path TEXT"),
            ("runs", "idempotency_key", "ALTER TABLE runs ADD COLUMN idempotency_key TEXT"),
            # BUG-06: persist ctx.trace JSON so /traces/{run_id} returns step-level data
            ("runs", "trace_json", "ALTER TABLE runs ADD COLUMN trace_json TEXT DEFAULT ''"),
            ("knowledge_bases", "ignore_dirs", "ALTER TABLE knowledge_bases ADD COLUMN ignore_dirs TEXT DEFAULT '[]'"),
            # Agent-KB integration (2026-06-02)
            ("agents", "kb_ids", "ALTER TABLE agents ADD COLUMN kb_ids TEXT DEFAULT '[]'"),
            ("agents", "kb_search_mode", "ALTER TABLE agents ADD COLUMN kb_search_mode TEXT DEFAULT 'bm25'"),
            # Run cancellation (2026-06-05): a non-zero value here is the
            # signal the cancel registry watches for to kill the active
            # subprocess + unwind the engine cleanly.
            ("runs", "cancel_requested", "ALTER TABLE runs ADD COLUMN cancel_requested INTEGER DEFAULT 0"),
            # Restart-resilient cancellation (2026-06-06): persist the live
            # claude subprocess PID so the cancel endpoint can fall back to
            # os.kill() even when the in-process cancel registry was lost
            # to a backend restart. Without this, run_5cbbb1d1dc54's orphan
            # subprocess kept writing for ~10 min after `agentpaas serve`
            # was re-launched, and the new instance's registry had no idea
            # the run existed.
            ("runs", "subprocess_pid", "ALTER TABLE runs ADD COLUMN subprocess_pid INTEGER"),
            # Per-agent working directory (2026-06-07): Bash tool's initial CWD
            # and the base directory for run workspace creation.
            ("agents", "work_dir", "ALTER TABLE agents ADD COLUMN work_dir TEXT DEFAULT ''"),
            # 3-directory system (2026-06-07):
            # source_dir: read-only input data; injected into system prompt preamble
            # run_dir:    workspace/run_{ts}/ base; traces & intermediate artifacts
            ("agents", "source_dir", "ALTER TABLE agents ADD COLUMN source_dir TEXT DEFAULT ''"),
            ("agents", "run_dir",    "ALTER TABLE agents ADD COLUMN run_dir TEXT DEFAULT ''"),
            # 会话记忆 P1（2026-06-12）：run 归属 thread（docs/MEMORY_DESIGN.md）
            ("runs", "thread_id", "ALTER TABLE runs ADD COLUMN thread_id TEXT"),
            # Cost tracking (2026-06-08): accurate USD cost from ClaudeCode JSONL.
            # cache_read/creation tokens are the dominant cost driver for heavy sessions
            # (reviewer agent: ~250K cache_read tokens per run ≈ $0.30–1.20).
            ("runs", "cost_usd",             "ALTER TABLE runs ADD COLUMN cost_usd REAL DEFAULT 0.0"),
            ("runs", "cache_read_tokens",    "ALTER TABLE runs ADD COLUMN cache_read_tokens INTEGER DEFAULT 0"),
            ("runs", "cache_creation_tokens","ALTER TABLE runs ADD COLUMN cache_creation_tokens INTEGER DEFAULT 0"),
        ]
        # Discovered while writing test_api_endpoints.py against a fresh
        # `:memory:` DB: `knowledge_bases` is created by `_ensure_table`
        # later in this function, so the probe and the ALTER both fire
        # against a non-existent table on first init → unhandled
        # OperationalError. Production hides this because the table
        # exists from a prior boot. Defensive fix: skip a migration when
        # its target table is missing — `_ensure_table` will lazy-create
        # the table next, with the new column already in its CREATE
        # statement.
        existing_tables = {
            row[0] for row in self.conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table'"
            ).fetchall()
        }
        for table, column, sql in migrations:
            if table not in existing_tables:
                continue
            try:
                self.conn.execute(f"SELECT {column} FROM {table} LIMIT 0")
            except sqlite3.OperationalError:
                self.conn.execute(sql)
        self.conn.commit()

        # Reap zombie jobs/runs left over from a prior crash or kill.
        # A row stuck in 'running' state for more than ZOMBIE_THRESHOLD must
        # belong to a process that no longer exists (since this Database
        # instance is constructed exactly once per process startup). Mark
        # them failed so the UI doesn't show "running" forever.
        # Without this, the workspace67 wiki compile from 9 days ago kept
        # appearing as "running" indefinitely. Threshold matches typical
        # safe completion windows: runs ≤ 6 h, kb_jobs ≤ 1 h.
        try:
            self.conn.execute(
                "UPDATE runs SET status='failed', "
                "  error=COALESCE(NULLIF(error,''),'') || "
                "    '[reaped at startup: process not running]', "
                "  completed_at=strftime('%Y-%m-%dT%H:%M:%S+00:00','now') "
                "WHERE status='running' "
                "  AND created_at < datetime('now','-6 hours')"
            )
            self.conn.execute(
                "UPDATE kb_jobs SET status='failed', "
                "  error=COALESCE(NULLIF(error,''),'') || "
                "    '[reaped at startup: process not running]', "
                "  completed_at=strftime('%Y-%m-%dT%H:%M:%S+00:00','now') "
                "WHERE status='running' "
                "  AND created_at < datetime('now','-1 hour')"
            )
            self.conn.commit()
        except sqlite3.OperationalError:
            # Tables may not exist yet on a brand-new DB (this migrate runs
            # before the kb tables are ensured below). Safe to skip — fresh
            # DB has no zombies by definition.
            pass

        # Knowledge base tables (added 2026-05-27)
        self._ensure_table("knowledge_bases", """
            CREATE TABLE IF NOT EXISTS knowledge_bases (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                name TEXT NOT NULL,
                root_dir TEXT NOT NULL,
                description TEXT DEFAULT '',
                created_at TEXT,
                updated_at TEXT
            )
        """)
        self._ensure_table("kb_files", """
            CREATE TABLE IF NOT EXISTS kb_files (
                id TEXT PRIMARY KEY,
                kb_id TEXT REFERENCES knowledge_bases(id),
                file_path TEXT NOT NULL,
                file_name TEXT NOT NULL,
                file_type TEXT,
                size INTEGER DEFAULT 0,
                status TEXT DEFAULT 'pending',
                added_at TEXT
            )
        """)
        self._ensure_table("kb_jobs", """
            CREATE TABLE IF NOT EXISTS kb_jobs (
                id TEXT PRIMARY KEY,
                kb_id TEXT REFERENCES knowledge_bases(id),
                tenant_id TEXT,
                job_type TEXT NOT NULL,
                status TEXT DEFAULT 'pending',
                progress INTEGER DEFAULT 0,
                total INTEGER DEFAULT 0,
                current_file TEXT DEFAULT '',
                error TEXT DEFAULT '',
                log TEXT DEFAULT '',
                started_at TEXT,
                completed_at TEXT,
                created_at TEXT
            )
        """)

        # Indexes for KB tables (FIX-06 extension)
        try:
            self.conn.executescript("""
                CREATE INDEX IF NOT EXISTS idx_kb_tenant ON knowledge_bases(tenant_id);
                CREATE INDEX IF NOT EXISTS idx_kbfiles_kb ON kb_files(kb_id);
                CREATE INDEX IF NOT EXISTS idx_kbjobs_kb ON kb_jobs(kb_id);
                CREATE INDEX IF NOT EXISTS idx_kbjobs_status ON kb_jobs(status);
                CREATE INDEX IF NOT EXISTS idx_kbjobs_tenant ON kb_jobs(tenant_id);
            """)
            self.conn.commit()
        except Exception:
            pass

    def execute(self, sql: str, params: tuple = ()) -> sqlite3.Cursor:
        return self.conn.execute(sql, params)

    def commit(self):
        self.conn.commit()

    def fetchone(self, sql: str, params: tuple = ()) -> Optional[dict]:
        row = self.conn.execute(sql, params).fetchone()
        return dict(row) if row else None

    def fetchall(self, sql: str, params: tuple = ()) -> List[dict]:
        return [dict(r) for r in self.conn.execute(sql, params).fetchall()]
