"""LLM client for the agent laboratory.

Design rules (CLAUDE.md, "LLM agent laboratory"):
* local OpenAI-compatible inference (mlx_lm.server on MacLustr nodes) is the default; optional BYOK providers are
  configured only through environment variables (never embedded credentials);
* every call is recorded: provider, model, system fingerprint, prompt and context hashes, full messages, temperature,
  seed, response, latency, token usage, estimated cost, cache hit;
* responses are cached content-addressed (sha256 of provider+model+messages+temperature+seed+schema) so a run is
  reproducible bit-for-bit from the cache;
* JSON actions are validated against a JSON schema; invalid output is retried once with a repair message and then
  rejected (the caller treats rejection as "no action" and logs it). LLM output is never accounting truth.
* `MockLLM` is a deterministic *policy surrogate* (rule-based) exposing the same interface; it is labelled as such in
  every log line and must never be reported as LLM inference.
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import time
from dataclasses import dataclass, field
from pathlib import Path

import httpx
import jsonschema

_JSON_RE = re.compile(r"\{.*\}", re.S)


def _hash(obj) -> str:
    return hashlib.sha256(json.dumps(obj, sort_keys=True, default=str).encode()).hexdigest()


def extract_json(text: str):
    """Return the first JSON object found in `text` (handles ```json fences), or None."""
    if text is None:
        return None
    t = text.strip()
    t = re.sub(r"^```(?:json)?\s*|\s*```$", "", t, flags=re.S)
    try:
        return json.loads(t)
    except Exception:
        pass
    m = _JSON_RE.search(t)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except Exception:
        return None


@dataclass
class ProviderConfig:
    name: str                      # 'local-qwen', 'anthropic', ...
    kind: str                      # 'openai_compat' | 'anthropic' | 'mock'
    model: str
    base_url: str = ''
    api_key_env: str = ''          # name of the env var holding the key (optional providers only)
    price_in_per_mtok: float = 0.0 # USD per million input tokens (0 for local)
    price_out_per_mtok: float = 0.0
    timeout: float = 120.0
    merge_system: bool = False     # chat templates that reject a system role (Mistral): prepend system text to the first user turn

    @property
    def api_key(self):
        return os.environ.get(self.api_key_env, '') if self.api_key_env else ''


LOCAL_PROVIDERS = {
    'qwen7b': ProviderConfig('qwen7b', 'openai_compat', 'mlx-community/Qwen2.5-7B-Instruct-4bit', 'http://localhost:8081/v1'),
    'mistral7b': ProviderConfig('mistral7b', 'openai_compat', 'mlx-community/Mistral-7B-Instruct-v0.3-4bit', 'http://localhost:8082/v1', merge_system=True),
    'llama3b': ProviderConfig('llama3b', 'openai_compat', 'mlx-community/Llama-3.2-3B-Instruct-4bit', 'http://localhost:8083/v1'),
    'mock': ProviderConfig('mock', 'mock', 'rule-surrogate-v1'),
}


@dataclass
class LLMResult:
    parsed: dict | None
    raw: str
    valid: bool
    cached: bool
    latency: float
    usage: dict
    cost_usd: float
    model: str
    fingerprint: str
    attempts: int
    error: str = ''


class LLMClient:
    def __init__(self, provider: ProviderConfig, cache_dir: str | Path = 'results/llm_cache',
                 audit_path: str | Path | None = None, run_id: str = 'run'):
        self.p = provider
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.audit_path = Path(audit_path) if audit_path else Path('results/llm_audit') / f'{run_id}.jsonl'
        self.audit_path.parent.mkdir(parents=True, exist_ok=True)
        self.run_id = run_id
        self.n_calls = 0
        self.n_cached = 0
        self.n_invalid = 0
        self.total_tokens = 0
        self.total_cost = 0.0
        self._http = httpx.Client(timeout=provider.timeout)

    # ------------------------------------------------------------------ raw completion
    def _complete_raw(self, messages, temperature, seed, max_tokens):
        p = self.p
        if p.kind == 'openai_compat':
            if p.merge_system and messages and messages[0]['role'] == 'system':
                sys_txt = messages[0]['content']
                rest = [dict(m) for m in messages[1:]]
                if rest and rest[0]['role'] == 'user':
                    rest[0]['content'] = sys_txt + '\n\n' + rest[0]['content']
                messages = rest
            headers = {'Content-Type': 'application/json'}
            if p.api_key:
                headers['Authorization'] = f'Bearer {p.api_key}'
            body = dict(model=p.model, messages=messages, temperature=temperature, max_tokens=max_tokens, seed=seed)
            r = self._http.post(f'{p.base_url}/chat/completions', json=body, headers=headers)
            r.raise_for_status()
            j = r.json()
            text = j['choices'][0]['message']['content']
            usage = j.get('usage', {})
            fp = j.get('system_fingerprint', '')
            return text, usage, fp
        if p.kind == 'anthropic':
            if not p.api_key:
                raise RuntimeError(f'missing credential in env {p.api_key_env}')
            headers = {'x-api-key': p.api_key, 'anthropic-version': '2023-06-01', 'content-type': 'application/json'}
            system = '\n'.join(m['content'] for m in messages if m['role'] == 'system')
            msgs = [m for m in messages if m['role'] != 'system']
            body = dict(model=p.model, max_tokens=max_tokens, temperature=temperature, messages=msgs)
            if system:
                body['system'] = system
            r = self._http.post(f'{p.base_url or "https://api.anthropic.com"}/v1/messages', json=body, headers=headers)
            r.raise_for_status()
            j = r.json()
            text = ''.join(b.get('text', '') for b in j.get('content', []))
            u = j.get('usage', {})
            usage = dict(prompt_tokens=u.get('input_tokens', 0), completion_tokens=u.get('output_tokens', 0))
            return text, usage, j.get('model', '')
        raise ValueError(p.kind)

    # ------------------------------------------------------------------ structured call
    def complete_json(self, messages: list[dict], schema: dict, temperature: float = 0.0, seed: int = 0,
                      max_tokens: int = 300, tags: dict | None = None) -> LLMResult:
        key = _hash(dict(provider=self.p.name, model=self.p.model, messages=messages, temperature=temperature,
                         seed=seed, schema=schema))
        cpath = self.cache_dir / f'{key}.json'
        cached = cpath.exists()
        t0 = time.time()
        attempts = 0
        error = ''
        if cached:
            rec = json.loads(cpath.read_text())
            text, usage, fp, attempts = rec['raw'], rec['usage'], rec['fingerprint'], rec['attempts']
            parsed = extract_json(text)
        else:
            msgs = list(messages)
            parsed, text, usage, fp = None, '', {}, ''
            for attempt in range(2):
                attempts += 1
                try:
                    text, usage, fp = self._complete_raw(msgs, temperature, seed, max_tokens)
                except Exception as e:                      # network / provider error
                    error = f'{type(e).__name__}: {e}'
                    text = ''
                    break
                parsed = extract_json(text)
                if parsed is not None:
                    try:
                        jsonschema.validate(parsed, schema)
                        break
                    except jsonschema.ValidationError as e:
                        error = f'schema: {e.message}'
                        parsed = None
                msgs = msgs + [{'role': 'assistant', 'content': text},
                               {'role': 'user', 'content': 'Your reply was not a valid JSON object matching the schema. '
                                                           'Reply with ONLY the JSON object, no prose.'}]
            if not error.startswith(('ConnectError', 'ReadTimeout', 'HTTPStatusError', 'RemoteProtocolError', 'ConnectTimeout')):
                cpath.write_text(json.dumps(dict(raw=text, usage=usage, fingerprint=fp, attempts=attempts, error=error)))
        valid = parsed is not None
        if valid:
            try:
                jsonschema.validate(parsed, schema)
            except jsonschema.ValidationError as e:
                valid, error = False, f'schema: {e.message}'
        latency = time.time() - t0
        ptok = int(usage.get('prompt_tokens', 0) or 0)
        ctok = int(usage.get('completion_tokens', 0) or 0)
        cost = (ptok * self.p.price_in_per_mtok + ctok * self.p.price_out_per_mtok) / 1e6
        self.n_calls += 1
        self.n_cached += int(cached)
        self.n_invalid += int(not valid)
        self.total_tokens += ptok + ctok
        self.total_cost += cost
        res = LLMResult(parsed if valid else None, text, valid, cached, latency, usage, cost, self.p.model, fp, attempts, error)
        self._audit(messages, schema, temperature, seed, res, tags or {})
        return res

    def _audit(self, messages, schema, temperature, seed, res: LLMResult, tags):
        rec = dict(ts=time.time(), run_id=self.run_id, provider=self.p.name, kind=self.p.kind, model=self.p.model,
                   fingerprint=res.fingerprint, surrogate=(self.p.kind == 'mock'),
                   prompt_hash=_hash(messages[0]['content'] if messages else ''), context_hash=_hash(messages),
                   schema_hash=_hash(schema), messages=messages, temperature=temperature, seed=seed,
                   response=res.raw, parsed=res.parsed, valid=res.valid, error=res.error, cached=res.cached,
                   latency=res.latency, usage=res.usage, cost_usd=res.cost_usd, attempts=res.attempts, tags=tags)
        with self.audit_path.open('a') as f:
            f.write(json.dumps(rec, default=str) + '\n')

    def stats(self):
        return dict(provider=self.p.name, model=self.p.model, calls=self.n_calls, cached=self.n_cached,
                    invalid=self.n_invalid, tokens=self.total_tokens, cost_usd=self.total_cost)


class MockLLM(LLMClient):
    """Deterministic rule-based *surrogate* with the LLM interface. Produces a mean-variance-like order from the
    observation JSON embedded in the last user message. Labelled surrogate=True in every audit record."""

    def __init__(self, **kw):
        super().__init__(LOCAL_PROVIDERS['mock'], **kw)

    def _complete_raw(self, messages, temperature, seed, max_tokens):
        obs = extract_json(messages[-1]['content']) or {}
        price = float(obs.get('price', 1.0))
        val = float(obs.get('own_valuation', price))
        h = float(obs.get('holdings', 0.0))
        cash = float(obs.get('cash', 0.0))
        target = float(obs.get('target_holdings', h))
        q = target - h
        side = 'buy' if q > 0 else ('sell' if q < 0 else 'hold')
        qty = abs(round(q, 2))
        limit = round(val * (0.995 if side == 'buy' else 1.005), 4)
        act = dict(side=side, quantity=qty, order_type='limit', limit_price=limit, borrow=0.0,
                   rationale='surrogate: move toward target holdings')
        return json.dumps(act), dict(prompt_tokens=0, completion_tokens=0), 'mock-surrogate'