"""Small HTTP transport and owned single-GPU native-server lifecycle."""
from contextlib import contextmanager
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import signal
import socket
import subprocess
import threading
import time
import urllib.request


class Client:
    def __init__(self, base='http://127.0.0.1:5991', adapter_id=None, timeout=120):
        self.base=base.rstrip('/').removesuffix('/v1')
        self.adapter_id=adapter_id
        self.timeout=timeout
        self.calls=[]
        self.lock=threading.RLock()

    def health(self):
        with urllib.request.urlopen(self.base+'/health',timeout=5) as stream:return json.load(stream)

    def verify_profile(self):
        with urllib.request.urlopen(self.base+'/lora-adapters',timeout=5) as stream:loaded=json.load(stream)
        if self.adapter_id is None and loaded:
            raise RuntimeError('Base-only profile requires a backend with no adapters loaded')
        if self.adapter_id is not None and self.adapter_id not in [r['id'] for r in loaded]:
            raise RuntimeError('Requested adapter ID is not loaded by the backend')
        return loaded

    def post(self, path, payload):
        body=dict(payload)
        if self.adapter_id is None:
            body.pop('lora',None)
        else:
            body['lora']=[dict(id=self.adapter_id,scale=1)]
        body.update(cache_prompt=False,chat_template_kwargs=dict(enable_thinking=False))
        req=urllib.request.Request(self.base+path,data=json.dumps(body).encode(),headers={'Content-Type':'application/json'})
        tick=time.perf_counter()
        with self.lock, urllib.request.urlopen(req,timeout=self.timeout) as stream:
            result=json.load(stream)
        usage=result.get('usage') or {}
        self.calls.append(dict(seconds=time.perf_counter()-tick,prompt_tokens=usage.get('prompt_tokens'),
                               completion_tokens=usage.get('completion_tokens'),
                               cached_tokens=(usage.get('prompt_tokens_details') or {}).get('cached_tokens',0),
                               finish_reason=(result.get('choices') or [{}])[0].get('finish_reason'),adapter_id=self.adapter_id))
        return result

    def generate(self, prompt, max_tokens=1536, json_output=False, schema=None):
        body=dict(model='qwen',messages=[dict(role='user',content=prompt)],max_tokens=max_tokens,temperature=0)
        if json_output: body['response_format']={'type':'json_object'}
        if schema is not None:body['response_format']={'type':'json_object','schema':schema}
        result=self.post('/v1/chat/completions',body)
        choice=result['choices'][0]
        if choice['finish_reason']=='length':
            raise ValueError('Structured response was truncated; no actions may execute')
        return choice['message'].get('content') or ''


class NativeServer:
    def __init__(self, root, gpu=0, port=5991, model='models/Ternary-Bonsai-2-27B-PQ2_0.gguf', adapter=None,
                 binary='llama.cpp-b2/build/bin/llama-server', context=4096):
        self.root=Path(root); self.root.mkdir(parents=True,exist_ok=True)
        self.gpu,self.port,self.model,self.adapter,self.binary,self.context=gpu,port,str(model),adapter,str(binary),context
        self.proc=None; self.samples=[]; self.stop=threading.Event()

    def event(self,kind,**values):
        row=dict(time_utc=datetime.now(timezone.utc).isoformat(),kind=kind,**values)
        with (self.root/'runtime-events.jsonl').open('a') as f:f.write(json.dumps(row)+'\n')

    def memory(self):
        return int(subprocess.check_output(['nvidia-smi',f'--id={self.gpu}','--query-gpu=memory.used','--format=csv,noheader,nounits'],text=True))

    def __enter__(self):
        with socket.socket() as sock:
            if sock.connect_ex(('127.0.0.1',self.port))==0:raise RuntimeError('Port occupied; no existing server will be stopped')
        if self.memory()>=100:raise RuntimeError(f'GPU {self.gpu} is occupied')
        env=dict(os.environ,CUDA_VISIBLE_DEVICES=str(self.gpu),PYTHONDONTWRITEBYTECODE='1')
        env['LD_LIBRARY_PATH']=str(Path(self.binary).resolve().parent)+((':'+env['LD_LIBRARY_PATH']) if env.get('LD_LIBRARY_PATH') else '')
        command=[self.binary,'-m',self.model,'--host','0.0.0.0','--port',str(self.port),'-c',str(self.context),'-ngl','99','-fa','on',
                 '--jinja','--reasoning','off','-np','1','--no-context-shift','--cache-ram','0','--no-cache-idle-slots']
        if self.adapter:command+=['--lora',str(self.adapter)]
        self.log=(self.root/f'server-{time.time_ns()}.log').open('x')
        self.proc=subprocess.Popen(command,env=env,stdout=self.log,stderr=subprocess.STDOUT)
        self.event('server_started',pid=self.proc.pid,gpu=self.gpu,command=command)
        try:
            deadline=time.monotonic()+180
            while time.monotonic()<deadline:
                if self.proc.poll() is not None:raise RuntimeError('Native server exited during startup')
                try:
                    with urllib.request.urlopen(f'http://127.0.0.1:{self.port}/health',timeout=1) as r:
                        if r.status==200:break
                except OSError:pass
                time.sleep(.25)
            else:raise TimeoutError('Native server readiness timed out')
            with urllib.request.urlopen(f'http://127.0.0.1:{self.port}/lora-adapters',timeout=5) as r:loaded=json.load(r)
            if [x['path'] for x in loaded] != ([str(self.adapter)] if self.adapter else []):
                raise RuntimeError('Loaded adapters do not match requested profile')
            def sample():
                while not self.stop.is_set():
                    self.samples.append(dict(time=time.time(),mib=self.memory()))
                    self.stop.wait(.5)
            self.thread=threading.Thread(target=sample,daemon=True);self.thread.start()
            self.event('server_ready',memory_mib=self.memory(),loaded_adapters=loaded)
            return Client(f'http://127.0.0.1:{self.port}',0 if self.adapter else None)
        except BaseException:
            self.__exit__(None,None,None)
            raise

    def __exit__(self,*args):
        self.stop.set()
        if hasattr(self,'thread'):self.thread.join(timeout=5)
        if self.proc:
            for sig in [signal.SIGTERM,signal.SIGINT,signal.SIGKILL]:
                if self.proc.poll() is not None:break
                self.proc.send_signal(sig)
                try:self.proc.wait(timeout=10)
                except subprocess.TimeoutExpired:pass
            self.log.close()
            self.event('server_stopped',pid=self.proc.pid,exit_code=self.proc.poll(),peak_mib=max((s['mib'] for s in self.samples),default=None))
            with (self.root/f'memory-{self.proc.pid}.json').open('x') as f:json.dump(self.samples,f)