"""Small SQLite repository. Queue transitions are serialized transactions."""
from contextlib import contextmanager
from datetime import datetime, timezone
import json
import sqlite3
import uuid

TERMINAL = {"succeeded", "failed", "cancelled", "timed_out", "interrupted"}


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


class Store:
    def __init__(self, settings):
        self.settings = settings
        settings.prepare()
        with self.connect() as db:
            version = db.execute("PRAGMA user_version").fetchone()[0]
            if version not in (0, 1):
                raise RuntimeError("Unsupported database version")
            db.executescript("""
                PRAGMA journal_mode=WAL;
                CREATE TABLE IF NOT EXISTS scripts(id TEXT PRIMARY KEY, data TEXT NOT NULL);
                CREATE TABLE IF NOT EXISTS versions(
                  script_id TEXT NOT NULL, version INTEGER NOT NULL, data TEXT NOT NULL,
                  PRIMARY KEY(script_id,version));
                CREATE TABLE IF NOT EXISTS runs(
                  id TEXT PRIMARY KEY, request_id TEXT NOT NULL UNIQUE,
                  script_id TEXT NOT NULL, state TEXT NOT NULL,
                  cancel_requested INTEGER NOT NULL DEFAULT 0, data TEXT NOT NULL);
                CREATE TABLE IF NOT EXISTS health(key TEXT PRIMARY KEY, value TEXT NOT NULL);
                PRAGMA user_version=1;
            """)

    @contextmanager
    def connect(self, write=False):
        db = sqlite3.connect(self.settings.db, timeout=10)
        db.row_factory = sqlite3.Row
        try:
            if write:
                db.execute("BEGIN IMMEDIATE")
            yield db
            db.commit()
        except Exception:
            db.rollback()
            raise
        finally:
            db.close()

    def scripts(self):
        with self.connect() as db:
            return [json.loads(r[0]) for r in db.execute("SELECT data FROM scripts ORDER BY id")]

    def script(self, identity, versions=False):
        with self.connect() as db:
            row = db.execute("SELECT data FROM scripts WHERE id=?", (identity,)).fetchone()
            if not row:
                raise KeyError("脚本不存在")
            item = json.loads(row[0])
            if versions:
                item["versions"] = [json.loads(r[0]) for r in db.execute(
                    "SELECT data FROM versions WHERE script_id=? ORDER BY version DESC", (identity,))]
            return item

    def save_script(self, item):
        with self.connect(write=True) as db:
            db.execute("INSERT OR REPLACE INTO scripts VALUES(?,?)", (item["id"], json.dumps(item)))

    def version(self, identity, version):
        with self.connect() as db:
            row = db.execute("SELECT data FROM versions WHERE script_id=? AND version=?", (identity, version)).fetchone()
            if not row:
                raise KeyError("版本不存在")
            return json.loads(row[0])

    def upload_version(self, item, content, filename):
        # A transaction prevents two simultaneous uploads from assigning the same version.
        with self.connect(write=True) as db:
            current = db.execute("SELECT data FROM scripts WHERE id=?", (item["id"],)).fetchone()
            if current:
                item = json.loads(current[0])
            version = item["version"] + 1
            target = self.settings.library / item["id"] / "versions" / str(version) / filename
            target.parent.mkdir(parents=True, exist_ok=False)
            try:
                with target.open("xb") as handle:
                    handle.write(content)
                import hashlib
                value = {"version": version, "sha256": hashlib.sha256(content).hexdigest(),
                         "created_at": now(), "filename": filename, "path": str(target)}
                item.update(version=version, sha256=value["sha256"], source_path=str(target), updated_at=now())
                db.execute("INSERT INTO versions VALUES(?,?,?)", (item["id"], version, json.dumps(value)))
                db.execute("INSERT OR REPLACE INTO scripts VALUES(?,?)", (item["id"], json.dumps(item)))
            except Exception:
                target.unlink(missing_ok=True)
                target.parent.rmdir()
                raise
        return item

    @staticmethod
    def unpack_run(row):
        data = json.loads(row["data"])
        data.update(state=row["state"], cancel_requested=bool(row["cancel_requested"]))
        return data

    def runs(self, script_id=None, limit=100):
        with self.connect() as db:
            query = "SELECT * FROM runs"
            values = []
            if script_id:
                query += " WHERE script_id=?"
                values.append(script_id)
            query += " ORDER BY rowid DESC LIMIT ?"
            values.append(limit)
            return [self.unpack_run(r) for r in db.execute(query, values)]

    def run(self, identity):
        with self.connect() as db:
            row = db.execute("SELECT * FROM runs WHERE id=?", (identity,)).fetchone()
            if not row:
                raise KeyError("运行记录不存在")
            return self.unpack_run(row)

    def run_for_request(self, request_id):
        with self.connect() as db:
            row = db.execute("SELECT * FROM runs WHERE request_id=?", (request_id,)).fetchone()
            return self.unpack_run(row) if row else None

    def queue(self, script, request_id, ip, execution=None):
        with self.connect(write=True) as db:
            old = db.execute("SELECT * FROM runs WHERE request_id=?", (request_id,)).fetchone()
            if old:
                if old["script_id"] != script["id"]:
                    raise ValueError("请求编号已用于其他脚本")
                return self.unpack_run(old)
            pending = db.execute("SELECT count(*) FROM runs WHERE state IN ('queued','running')").fetchone()[0]
            if pending >= 20:
                raise ValueError("队列已满，请等待已有测试结束")
            item = {"id": uuid.uuid4().hex, "script_id": script["id"], "script_name": script["name"],
                    "version": script["version"], "sha256": script["sha256"], "state": "queued",
                    "created_at": now(), "started_at": None, "finished_at": None, "exit_code": None,
                    "message": "等待执行", "config": script["config"], "artifacts": [],
                    "synthetic": script.get("synthetic", False), "source_ip": ip,
                    "source_path": script["source_path"], "language": script["language"],
                    "kind": script["kind"], "source_group": script["source_group"]}
            if execution:
                item["execution"] = execution
            db.execute("INSERT INTO runs(id,request_id,script_id,state,data) VALUES(?,?,?,?,?)",
                       (item["id"], request_id, script["id"], "queued", json.dumps(item)))
            return item

    def claim(self):
        with self.connect(write=True) as db:
            if db.execute("SELECT 1 FROM runs WHERE state='running'").fetchone():
                return None
            row = db.execute("SELECT * FROM runs WHERE state='queued' ORDER BY rowid LIMIT 1").fetchone()
            if not row:
                return None
            item = self.unpack_run(row)
            item.update(started_at=now(), state="running", message="正在启动测试")
            db.execute("UPDATE runs SET state='running',data=? WHERE id=?", (json.dumps(item), item["id"]))
            return item

    def finish(self, identity, state, **values):
        assert state in TERMINAL
        with self.connect(write=True) as db:
            row = db.execute("SELECT * FROM runs WHERE id=?", (identity,)).fetchone()
            item = self.unpack_run(row)
            item.update(values, state=state, finished_at=now())
            db.execute("UPDATE runs SET state=?,data=? WHERE id=?", (state, json.dumps(item), identity))
            return item

    def running_message(self, identity, message):
        with self.connect(write=True) as db:
            row = db.execute("SELECT * FROM runs WHERE id=? AND state='running'", (identity,)).fetchone()
            if row:
                item = self.unpack_run(row)
                item["message"] = message
                db.execute("UPDATE runs SET data=? WHERE id=?", (json.dumps(item), identity))

    def update_execution(self, identity, values):
        with self.connect(write=True) as db:
            row = db.execute("SELECT * FROM runs WHERE id=? AND state='running'", (identity,)).fetchone()
            if not row:
                raise ValueError("运行状态已变化")
            item = self.unpack_run(row)
            item["execution"] = {**item.get("execution", {}), **values}
            db.execute("UPDATE runs SET data=? WHERE id=?", (json.dumps(item), identity))
            return item

    def cancel(self, identity):
        with self.connect(write=True) as db:
            row = db.execute("SELECT * FROM runs WHERE id=?", (identity,)).fetchone()
            if not row:
                raise KeyError("运行记录不存在")
            item = self.unpack_run(row)
            if row["state"] == "queued":
                item.update(state="cancelled", finished_at=now(), message="已取消排队")
                db.execute("UPDATE runs SET state='cancelled',cancel_requested=1,data=? WHERE id=?",
                           (json.dumps(item), identity))
            elif row["state"] == "running":
                db.execute("UPDATE runs SET cancel_requested=1 WHERE id=?", (identity,))
        return self.run(identity)

    def heartbeat(self):
        with self.connect(write=True) as db:
            db.execute("INSERT OR REPLACE INTO health VALUES('worker',?)", (now(),))

    def health(self):
        with self.connect() as db:
            row = db.execute("SELECT value FROM health WHERE key='worker'").fetchone()
        last = row[0] if row else None
        online = bool(last and (datetime.now(timezone.utc) - datetime.fromisoformat(last)).total_seconds() < 15)
        return {"online": online, "last_seen": last}