"""Optional trace composition. API authentication is independent of Web auth."""
from contextlib import asynccontextmanager
import asyncio
import hmac
import json
import os

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
from contextra_trace import set_recorder
from contextra_trace.http import Capture
from contextra_trace.writer import Recorder
from app.trace_store import Store


def install(app, token, server, manager=None):
    path = os.environ.get('CONTEXTRA_TRACE_POLICY_FILE')
    if not path:
        return
    local = os.environ.get('CONTEXTRA_TRACE_SERVER', server)
    store = Store(local)
    recorder = Recorder(store.write, server=local, path=path, start=False)
    set_recorder(recorder)
    last = {}
    def observe():
        if manager is None or not manager._lock.acquire(blocking=False):
            return
        try:
            groups = list(manager._groups.values())[:512]
        finally:
            manager._lock.release()
        seen = set()
        for group in groups:
            lifecycle = group.lifecycle
            if not lifecycle._lock.acquire(blocking=False):
                continue
            try:
                state, run = lifecycle.observed_state.value, lifecycle.run_id
                error = lifecycle.last_error.as_dict() if lifecycle.last_error else None
            finally:
                lifecycle._lock.release()
            key = str(group.stream_id)
            seen.add(key)
            current = (run,state,json.dumps(error,sort_keys=True))
            if last.get(key) != current:
                recorder.emit(dict(component='runtime',phase='STATE_OBSERVED',
                    cameraId=str(group.spec.get('camera_id') or group.spec.get('source_id') or key),
                    runId=run,observedState=state,error=error))
                last[key] = current
        for key in set(last)-seen:
            del last[key]
    recorder.observer = observe

    @asynccontextmanager
    async def lifetime(_):
        import threading
        recorder.thread = threading.Thread(target=recorder.run, name='contextra-trace', daemon=True)
        recorder.thread.start()
        try:
            # Empty write initializes the diagnostic tables without blocking readiness.
            recorder.emit({'component':'runtime','phase':'TRACE_STARTED'})
            yield
        finally:
            await asyncio.to_thread(recorder.close)
            set_recorder(None)

    router = APIRouter(lifespan=lifetime)

    @router.post('/internal/trace/{action}')
    async def manage(action: str, request: Request):
        auth = request.headers.getlist('authorization')
        if len(auth) != 1 or not hmac.compare_digest(auth[0].encode(), ('Bearer '+token).encode()):
            return JSONResponse({'error':'TRACE_UNAUTHORIZED'},status_code=401)
        raw = bytearray()
        try:
            async with asyncio.timeout(2):
                async for chunk in request.stream():
                    raw.extend(chunk)
                    if len(raw) > (4*1024**2 if action == 'ingest' else 8192):
                        raise ValueError('TRACE_BODY_LIMIT')
            value = json.loads(raw)
            if action == 'ingest':
                if not isinstance(value,list) or len(value)>32 or any(not isinstance(d,dict) for d in value):
                    raise ValueError('TRACE_BATCH_INVALID')
                accepted = sum(recorder.emit(dict(d,sourceTraceId=d.get('traceId'),
                                                  sourceRecordedAtMs=d.get('recordedAtMs'))) for d in value)
                return JSONResponse({'queued':accepted,'dropped':len(value)-accepted},status_code=202)
            if not isinstance(value,dict):
                raise ValueError('TRACE_REQUEST_INVALID')
            if value.get('server') != local:
                raise ValueError('TRACE_SERVER_MISMATCH')
            if action == 'status':
                result = recorder.status()
                try:
                    result.update(await asyncio.to_thread(store.query,action,value))
                except Exception:
                    result['database']='TRACE_DB_UNAVAILABLE'
                return result
            return await asyncio.to_thread(store.query,action,value)
        except (ValueError,TypeError,TimeoutError):
            return JSONResponse({'error':'TRACE_REQUEST_INVALID'},status_code=422)
        except Exception:
            return JSONResponse({'error':'TRACE_DB_UNAVAILABLE'},status_code=503)

    app.include_router(router)
    app.add_middleware(Capture, recorder=recorder, component='runtime')
    app.state.trace = recorder