"""EVART-02: durable per-channel claims and globally bounded reservations."""
from contextlib import contextmanager
import json
import shutil
import threading
import time
import uuid

from app.event_artifact_contract import RESERVATION, relative_directory


class ArtifactStore:
    def __init__(self, config):
        self.config = config
        self.local = threading.local()
        self.connections = []
        self.connection_lock = threading.Lock()

    @contextmanager
    def transaction(self):
        import psycopg2
        from psycopg2.extras import RealDictCursor
        conn = getattr(self.local, 'connection', None)
        if conn is None or conn.closed:
            conn = psycopg2.connect(self.config.database, connect_timeout=3,
                                   options='-c statement_timeout=3000 -c lock_timeout=2000')
            self.local.connection = conn
            with self.connection_lock:
                self.connections.append(conn)
        with conn:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                yield cursor

    def close(self):
        with self.connection_lock:
            for conn in self.connections:
                conn.close()
            self.connections.clear()

    def initialize(self):
        self.config.root.mkdir(parents=True, exist_ok=True)
        (self.config.root / '.staging').mkdir(exist_ok=True)
        from app.event_artifact_schema import initialize
        with self.transaction() as cur:
            initialize(cur)

    def recover(self):
        # Parent holds .worker.lock and all previous workers have exited.
        with self.transaction() as cur:
            cur.execute('SELECT id FROM contextra_artifact_storage WHERE id=1 FOR UPDATE')
            cur.execute("UPDATE contextra_event_artifacts SET artifact_state='FAILED',error='PROCESS_INTERRUPTED' WHERE artifact_state IN ('RENDERING','CAPTURING')")
            cur.execute("UPDATE contextra_event_artifacts SET delivery_state=CASE WHEN attempts<2 THEN 'RETRY' ELSE 'FAILED' END,delivery_error='RESPONSE_UNKNOWN' WHERE delivery_state='SENDING'")
            cur.execute("UPDATE contextra_event_artifacts SET delivery_state=CASE WHEN attempts=0 THEN 'PENDING' ELSE 'RETRY' END WHERE delivery_state='CLAIMED'")
            cur.execute('UPDATE contextra_event_artifacts SET claim_token=NULL WHERE claim_token IS NOT NULL')
            cur.execute('''INSERT INTO contextra_artifact_channels(channel_key)
                SELECT DISTINCT channel_key FROM contextra_event_artifacts ON CONFLICT DO NOTHING''')
            cur.execute('''UPDATE contextra_artifact_channels c SET render_token=NULL,send_token=NULL,
                reserved_jobs=(SELECT count(*) FROM contextra_event_artifacts e
                WHERE e.channel_key=c.channel_key AND reservation>0)''')
            cur.execute('''UPDATE contextra_artifact_storage SET
                reserved_jobs=(SELECT count(*) FROM contextra_event_artifacts WHERE reservation>0),
                reserved_bytes=(SELECT COALESCE(sum(reservation),0) FROM contextra_event_artifacts) WHERE id=1''')
        self.scan_storage()
        self.expire()

    def scan_storage(self):
        with self.transaction() as cur:
            cur.execute('SELECT accounting_revision FROM contextra_artifact_storage WHERE id=1')
            revision = cur.fetchone()['accounting_revision']
        size = 0
        for path in self.config.root.rglob('*'):
            try:
                if path.is_file() and not path.is_symlink():
                    size += path.stat().st_size
            except FileNotFoundError:
                pass  # A completed job atomically moved its staging directory.
        with self.transaction() as cur:
            cur.execute('UPDATE contextra_artifact_storage SET observed_bytes=%s,observed_at=%s WHERE id=1 AND accounting_revision=%s', (size,time.time(),revision))
        return size

    def accept(self, event):
        relative = relative_directory(event)
        channel = event.get('streamId', str(event['cameraId']))
        confirmed = event['aiRecordedAtWallMs']/1000
        victim = None
        with self.transaction() as cur:
            cur.execute('SELECT * FROM contextra_artifact_storage WHERE id=1 FOR UPDATE')
            storage = cur.fetchone()
            cur.execute('SELECT 1 FROM contextra_event_artifacts WHERE event_id=%s', (event['eventId'],))
            if cur.fetchone():
                return False, 'DUPLICATE'
            cur.execute('INSERT INTO contextra_artifact_channels(channel_key) VALUES (%s) ON CONFLICT DO NOTHING', (channel,))
            cur.execute('SELECT reserved_jobs FROM contextra_artifact_channels WHERE channel_key=%s FOR UPDATE', (channel,))
            channel_full = cur.fetchone()['reserved_jobs'] >= self.config.channel_limit
            if confirmed+self.config.lifetime > time.time() and (channel_full or storage['reserved_jobs'] >= 64):
                cur.execute('''SELECT * FROM contextra_event_artifacts WHERE reservation>0
                    AND artifact_state IN ('CAPTURE_READY','EXPIRED') AND claim_token IS NULL
                    AND (%s=false OR channel_key=%s)
                    ORDER BY (confirmed_at>%s),confirmed_at,event_id LIMIT 1 FOR UPDATE SKIP LOCKED''',
                    (channel_full,channel,time.time()-self.config.lifetime))
                victim = cur.fetchone()
                if victim:
                    victim = dict(victim,claim_token=uuid.uuid4().hex,error='EVICTED_CAPACITY')
                    cur.execute("UPDATE contextra_event_artifacts SET artifact_state='EXPIRED',delivery_state='EXPIRED',error='EVICTED_CAPACITY',claim_token=%s WHERE event_id=%s",
                                (victim['claim_token'],victim['event_id']))
        if victim:
            self.cleanup_expired(victim)
        now, free = time.time(), shutil.disk_usage(self.config.root).free
        with self.transaction() as cur:
            cur.execute('SELECT * FROM contextra_artifact_storage WHERE id=1 FOR UPDATE')
            storage = cur.fetchone()
            cur.execute('SELECT 1 FROM contextra_event_artifacts WHERE event_id=%s', (event['eventId'],))
            if cur.fetchone():
                return False, 'DUPLICATE'
            cur.execute('SELECT reserved_jobs FROM contextra_artifact_channels WHERE channel_key=%s FOR UPDATE', (channel,))
            count = cur.fetchone()['reserved_jobs']
            available = free-storage['reserved_bytes']-RESERVATION-self.config.external_reserve
            paused = available < self.config.low_free or (storage['paused'] and available < self.config.resume_free)
            reason = ('EXPIRED_TIME' if confirmed+self.config.lifetime <= now else
                      'STORAGE_FREE_LOW' if paused else
                      'ARTIFACT_CAP' if storage['observed_bytes']+storage['reserved_bytes']+RESERVATION > self.config.cap else
                      'PENDING_CAP_PROTECTED' if storage['reserved_jobs'] >= 64 or count >= self.config.channel_limit else None)
            expired = reason == 'EXPIRED_TIME'
            cur.execute('UPDATE contextra_artifact_storage SET paused=%s WHERE id=1', (paused,))
            cur.execute('''INSERT INTO contextra_event_artifacts
                (event_id,payload,relative_path,artifact_state,delivery_state,reservation,error,created_at,updated_at,channel_key,confirmed_at)
                VALUES (%s,%s::jsonb,%s,%s,%s,%s,%s,%s,%s,%s,%s)''',
                (event['eventId'],json.dumps(event),relative,'EXPIRED' if expired else 'FAILED' if reason else 'CAPTURING',
                 'EXPIRED' if expired else 'PENDING',0 if reason else RESERVATION,reason,now,now,channel,confirmed))
            if not reason:
                cur.execute('UPDATE contextra_artifact_storage SET reserved_jobs=reserved_jobs+1,reserved_bytes=reserved_bytes+%s,accounting_revision=accounting_revision+1 WHERE id=1', (RESERVATION,))
                cur.execute('UPDATE contextra_artifact_channels SET reserved_jobs=reserved_jobs+1 WHERE channel_key=%s', (channel,))
        return True, reason

    def captured(self, event_id, error=None):
        with self.transaction() as cur:
            cur.execute('''UPDATE contextra_event_artifacts SET artifact_state=%s,error=%s,updated_at=%s
                WHERE event_id=%s AND artifact_state='CAPTURING' ''',
                ('FAILED' if error else 'CAPTURE_READY',error,time.time(),event_id))

    def expire(self):
        now = time.time()
        with self.transaction() as cur:
            cur.execute('''UPDATE contextra_event_artifacts SET
                artifact_state=CASE WHEN artifact_state='READY' THEN 'READY' ELSE 'EXPIRED' END,
                delivery_state='EXPIRED',error=COALESCE(error,'EXPIRED_TIME'),updated_at=%s
                WHERE confirmed_at<=%s AND claim_token IS NULL
                AND delivery_state IN ('PENDING','RETRY') AND artifact_state<>'CAPTURING' ''',
                (now,now-self.config.lifetime))

    def claim(self, kind):
        now, token = time.time(), uuid.uuid4().hex
        render = kind == 'render'
        column, served = ('render_token','last_render_at') if render else ('send_token','last_send_at')
        condition = ("e.reservation>0 AND e.artifact_state IN ('CAPTURE_READY','FAILED','EXPIRED')" if render else
                     "e.reservation=0 AND e.artifact_state IN ('READY','FAILED') AND e.delivery_state IN ('PENDING','RETRY') AND e.next_attempt_at<=%s AND e.confirmed_at>%s")
        args = [] if render else [now,now-self.config.lifetime]
        with self.transaction() as cur:
            cur.execute('SELECT id FROM contextra_artifact_storage WHERE id=1 FOR UPDATE')
            cur.execute(f'''SELECT e.* FROM contextra_event_artifacts e
                JOIN contextra_artifact_channels c USING(channel_key)
                WHERE {condition} AND e.claim_token IS NULL AND c.{column} IS NULL
                ORDER BY (e.confirmed_at<%s),c.{served},e.confirmed_at,e.event_id
                LIMIT 1 FOR UPDATE OF c,e SKIP LOCKED''', (*args,now-self.config.recent))
            row = cur.fetchone()
            if row is None:
                return None
            row = dict(row,claim_token=token)
            cur.execute(f'UPDATE contextra_artifact_channels SET {column}=%s,{served}=%s WHERE channel_key=%s',
                        (token,now,row['channel_key']))
            if render:
                cur.execute("UPDATE contextra_event_artifacts SET claim_token=%s,artifact_state=CASE WHEN artifact_state='CAPTURE_READY' THEN 'RENDERING' ELSE artifact_state END WHERE event_id=%s", (token,row['event_id']))
            else:
                cur.execute("UPDATE contextra_event_artifacts SET claim_token=%s,delivery_state='CLAIMED' WHERE event_id=%s", (token,row['event_id']))
            return row

    def release_claim(self, row, kind):
        column = 'render_token' if kind == 'render' else 'send_token'
        with self.transaction() as cur:
            cur.execute('SELECT id FROM contextra_artifact_storage WHERE id=1 FOR UPDATE')
            cur.execute('UPDATE contextra_event_artifacts SET claim_token=NULL WHERE event_id=%s AND claim_token=%s', (row['event_id'],row['claim_token']))
            cur.execute(f'UPDATE contextra_artifact_channels SET {column}=NULL WHERE channel_key=%s AND {column}=%s', (row['channel_key'],row['claim_token']))

    def cleanup_expired(self, row):
        stage = self.config.root/'.staging'/row['event_id']
        if stage.exists():
            if stage.is_symlink():
                raise ValueError('Unsafe staging path')
            shutil.rmtree(stage)
        self.rendered(row,{},row.get('error') or 'EXPIRED_TIME',0,expired=True)
        self.release_claim(row,'render')

    def rendered(self, row, media, error, size, expired=False):
        with self.transaction() as cur:
            cur.execute('SELECT id FROM contextra_artifact_storage WHERE id=1 FOR UPDATE')
            cur.execute('SELECT reservation,stored_bytes FROM contextra_event_artifacts WHERE event_id=%s AND claim_token=%s FOR UPDATE', (row['event_id'],row['claim_token']))
            previous = cur.fetchone()
            if previous is None:
                raise RuntimeError('Stale artifact completion')
            cur.execute('''UPDATE contextra_event_artifacts SET artifact_state=%s,
                delivery_state=CASE WHEN %s THEN 'EXPIRED' ELSE delivery_state END,
                media=%s::jsonb,error=%s,stored_bytes=%s,reservation=0,updated_at=%s WHERE event_id=%s''',
                ('EXPIRED' if expired else 'FAILED' if error else 'READY',expired,json.dumps(media),error,size,time.time(),row['event_id']))
            if previous['reservation']:
                cur.execute('UPDATE contextra_artifact_storage SET reserved_jobs=reserved_jobs-1,reserved_bytes=reserved_bytes-%s WHERE id=1', (previous['reservation'],))
                cur.execute('UPDATE contextra_artifact_channels SET reserved_jobs=reserved_jobs-1 WHERE channel_key=%s', (row['channel_key'],))
            cur.execute('UPDATE contextra_artifact_storage SET observed_bytes=GREATEST(0,observed_bytes+%s),accounting_revision=accounting_revision+1 WHERE id=1', (size-previous['stored_bytes'],))

    def prepare_delivery(self, row, body):
        allowed = self.config.delivery_cameras
        state = ('EXPIRED' if row['confirmed_at']+self.config.lifetime<=time.time() else
                 'SKIPPED' if allowed is not None and str(row['payload']['cameraId']) not in allowed else
                 'BLOCKED_TIME' if body is None else 'CLAIMED')
        with self.transaction() as cur:
            cur.execute('''UPDATE contextra_event_artifacts SET body=COALESCE(body,%s::jsonb),
                delivery_state=%s,delivery_error=%s,updated_at=%s
                WHERE event_id=%s AND claim_token=%s AND delivery_state='CLAIMED' RETURNING body''',
                (json.dumps(body),state,'TEST_DELIVERY_EXCLUDED' if state=='SKIPPED' else state if state!='CLAIMED' else None,
                 time.time(),row['event_id'],row['claim_token']))
            result = cur.fetchone()
            return result['body'] if result and state=='CLAIMED' else None

    def begin_attempt(self, row):
        with self.transaction() as cur:
            cur.execute('''UPDATE contextra_event_artifacts SET attempts=attempts+1,delivery_state='SENDING',updated_at=%s
                WHERE event_id=%s AND claim_token=%s AND attempts<2 AND delivery_state='CLAIMED'
                AND confirmed_at>%s RETURNING attempts''',
                (time.time(),row['event_id'],row['claim_token'],time.time()-self.config.lifetime))
            result = cur.fetchone()
            if not result:
                cur.execute("UPDATE contextra_event_artifacts SET delivery_state='EXPIRED',delivery_error='EXPIRED_TIME' WHERE event_id=%s AND claim_token=%s AND delivery_state='CLAIMED'", (row['event_id'],row['claim_token']))
            return result['attempts'] if result else None

    def finish_attempt(self, row, state, status, response, error):
        with self.transaction() as cur:
            cur.execute('''UPDATE contextra_event_artifacts SET delivery_state=%s,http_status=%s,response=%s::jsonb,
                delivery_error=%s,next_attempt_at=%s,updated_at=%s WHERE event_id=%s AND claim_token=%s''',
                (state,status,json.dumps(response),error,time.time()+1,time.time(),row['event_id'],row['claim_token']))

    def summary(self):
        with self.transaction() as cur:
            cur.execute('SELECT * FROM contextra_artifact_storage WHERE id=1')
            storage = dict(cur.fetchone())
            cur.execute('''SELECT channel_key,artifact_state,delivery_state,count(*) AS count,
                min(confirmed_at) AS oldest FROM contextra_event_artifacts
                WHERE reservation>0 OR delivery_state IN ('PENDING','RETRY','CLAIMED','SENDING')
                GROUP BY channel_key,artifact_state,delivery_state''')
            return dict(storage=storage,pending=[dict(r) for r in cur.fetchall()])