"""A local stand-in for the Claude Messages API (POST /v1/messages), for checking lesson code without
spending money or needing a real key.

  python scripts/tutorials/llm_stub.py [port]          # default 4010
  then run code with  ANTHROPIC_API_KEY=test  ANTHROPIC_BASE_URL=http://127.0.0.1:4010

It VALIDATES requests the way the real API would for common mistakes (model and max_tokens required,
messages must alternate correctly, every tool_result must answer a tool_use id from the previous turn,
JSON-schema formats must be objects) and answers deterministically:
  - plain prompts                → a short, sensible text reply (canned by topic)
  - output_config.format schema  → a JSON object that satisfies the schema
  - tools, no tool_result yet    → a tool_use block for the most relevant tool
  - after tool_result            → a final answer that quotes the tool result
  - "stream": true               → server-sent events (message_start … message_stop)
Real model answers are worded differently — lessons say so wherever they show stub output.
"""
import json
import re
import sys
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 4010
counter = {'n': 0}


def last_user_text(messages):
    for m in reversed(messages):
        if m['role'] == 'user':
            c = m['content']
            if isinstance(c, str):
                return c
            return ' '.join(b.get('text', '') for b in c if isinstance(b, dict) and b.get('type') == 'text')
    return ''


def canned_text(prompt, system):
    p = prompt.lower()
    if 'explain' in p and 'recursion' in p:
        return ('Recursion is when a function solves a problem by calling itself on a smaller version of the same '
                'problem, until it reaches a simple case it can answer directly (the base case).')
    if 'context:' in p or '<context>' in p or 'documents' in p:
        m = re.search(r'\[(\d+)\][^\n]*?(\d+\s*%[^\n.]*)', prompt)
        if 'attendance' in p:
            return 'You need at least 75% attendance in each subject to sit the end-semester exam [1].'
        return 'Based on the provided documents: ' + (m.group(0) if m else 'the answer is in the context.')
    if 'summar' in p:
        return 'Summary: the team shipped the login feature, found two bugs in payments, and will fix them before Friday.'
    cats = re.search(r'one of:\s*([a-z ,]+)\.', p)
    if cats:   # classification with a listed label set: pick by simple keywords (a real model reasons instead)
        labels = [c.strip() for c in cats.group(1).split(',') if c.strip()]
        message = p.rsplit('message:', 1)[-1]
        hints = {'exams': ['exam', 'result', 'marks', 'hall ticket'], 'fees': ['fee', 'pay', 'refund'],
                 'attendance': ['absent', 'attendance', 'bunk']}
        for label in labels:
            if label in message or any(h in message for h in hints.get(label, [])):
                return label
        return labels[-1]
    if 'sentiment' in p:
        return 'negative'
    if 'name' in p and 'capital' in p:
        return 'New Delhi'
    if system and 'tutor' in system.lower():
        return 'Good question! Think about what changes each time the loop runs — what should the loop stop on?'
    return 'Here is a short answer to your question.'


def fake_value(schema, key=''):
    t = schema.get('type')
    if 'enum' in schema:
        return schema['enum'][0]
    if t == 'object':
        return {k: fake_value(v, k) for k, v in schema.get('properties', {}).items()}
    if t == 'array':
        item = schema.get('items', {})
        if key == 'skills':
            return ['Java', 'Spring Boot', 'SQL']
        return [fake_value(item, key)]
    if t == 'integer':
        return {'experience_years': 2, 'graduation_year': 2025}.get(key, 1)
    if t == 'number':
        return {'cgpa': 8.4}.get(key, 1.0)
    if t == 'boolean':
        return True
    return {'name': 'Priya Sharma', 'email': 'priya.sharma@example.com', 'degree': 'B.Tech CSE',
            'college': 'Example Institute of Technology', 'summary': 'Backend developer intern with Java experience.'}.get(key, 'text')


def validate(body):
    for f in ('model', 'max_tokens', 'messages'):
        if f not in body:
            return f'{f}: Field required'
    msgs = body['messages']
    if not msgs or msgs[0]['role'] != 'user':
        return 'messages: first message must use the "user" role'
    for i, m in enumerate(msgs):
        c = m['content']
        if isinstance(c, list):
            for b in c:
                if b.get('type') == 'tool_result':
                    prev = msgs[i - 1] if i else None
                    ids = [x.get('id') for x in (prev['content'] if prev and isinstance(prev['content'], list) else []) if isinstance(x, dict) and x.get('type') == 'tool_use']
                    if b.get('tool_use_id') not in ids:
                        return f'messages.{i}: tool_result {b.get("tool_use_id")} has no matching tool_use in the previous message'
    fmt = (body.get('output_config') or {}).get('format')
    if fmt and fmt.get('type') == 'json_schema' and fmt.get('schema', {}).get('type') != 'object':
        return 'output_config.format.schema: must be an object schema'
    return None


def choose_tool(tools, prompt):
    words = set(re.findall(r'[a-z]+', prompt.lower()))
    def overlap(t):
        return len(words & set(re.findall(r'[a-z]+', (t['name'].replace('_', ' ') + ' ' + t.get('description', '')).lower())))
    return max(tools, key=overlap)


def tool_input_for(tool, prompt):
    props = tool.get('input_schema', {}).get('properties', {})
    out = {}
    for k, v in props.items():
        nums = re.findall(r'\d+', prompt)
        if v.get('type') in ('integer', 'number') and nums:
            out[k] = int(nums[0])
        elif 'city' in k:
            m = re.search(r'\b(Jaipur|Pune|Delhi|Mumbai|Bengaluru|Chennai)\b', prompt)
            out[k] = m.group(1) if m else 'Jaipur'
        else:
            out[k] = fake_value(v, k)
    return out


def respond(body):
    counter['n'] += 1
    msg_id = f'msg_stub_{counter["n"]:04d}'
    msgs = body['messages']
    system = body.get('system')
    if isinstance(system, list):
        system = ' '.join(b.get('text', '') for b in system)
    prompt = last_user_text(msgs)
    last = msgs[-1]['content']
    has_result = isinstance(last, list) and any(isinstance(b, dict) and b.get('type') == 'tool_result' for b in last)
    fmt = (body.get('output_config') or {}).get('format')

    if body.get('tools') and not has_result:
        tool = choose_tool(body['tools'], prompt)
        content = [{'type': 'tool_use', 'id': f'toolu_stub_{counter["n"]:04d}', 'name': tool['name'], 'input': tool_input_for(tool, prompt)}]
        stop = 'tool_use'
    elif has_result:
        results = [b for b in last if b.get('type') == 'tool_result']
        texts = []
        for r in results:
            c = r.get('content')
            texts.append(c if isinstance(c, str) else json.dumps(c))
        content = [{'type': 'text', 'text': 'Using the tool result: ' + '; '.join(texts)}]
        stop = 'end_turn'
    elif fmt and fmt.get('type') == 'json_schema':
        content = [{'type': 'text', 'text': json.dumps(fake_value(fmt['schema']))}]
        stop = 'end_turn'
    else:
        content = [{'type': 'text', 'text': canned_text(prompt, system)}]
        stop = 'end_turn'

    in_tokens = max(1, len(json.dumps(body)) // 4)
    out_tokens = max(1, len(json.dumps(content)) // 4)
    cached = 0
    sys_blocks = body.get('system') if isinstance(body.get('system'), list) else []
    if any('cache_control' in b for b in sys_blocks) and counter['n'] > 1:
        cached = sum(len(b.get('text', '')) for b in sys_blocks) // 4
    return {
        'id': msg_id, 'type': 'message', 'role': 'assistant', 'model': body['model'], 'content': content,
        'stop_reason': stop, 'stop_sequence': None,
        'usage': {'input_tokens': in_tokens - cached, 'output_tokens': out_tokens,
                  'cache_read_input_tokens': cached, 'cache_creation_input_tokens': 0},
    }


class Handler(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def do_POST(self):
        if not self.path.startswith('/v1/messages'):
            self.send_error(404)
            return
        body = json.loads(self.rfile.read(int(self.headers.get('content-length', 0))) or b'{}')
        err = validate(body)
        if err:
            payload = json.dumps({'type': 'error', 'error': {'type': 'invalid_request_error', 'message': err}}).encode()
            self.send_response(400)
            self.send_header('content-type', 'application/json')
            self.end_headers()
            self.wfile.write(payload)
            return
        msg = respond(body)
        if body.get('stream'):
            self.send_response(200)
            self.send_header('content-type', 'text/event-stream')
            self.end_headers()

            def ev(name, data):
                self.wfile.write(f'event: {name}\ndata: {json.dumps(data)}\n\n'.encode())

            start = dict(msg, content=[], stop_reason=None)
            ev('message_start', {'type': 'message_start', 'message': start})
            text = msg['content'][0].get('text', '')
            ev('content_block_start', {'type': 'content_block_start', 'index': 0, 'content_block': {'type': 'text', 'text': ''}})
            for piece in re.findall(r'\S+\s*', text):
                ev('content_block_delta', {'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': piece}})
            ev('content_block_stop', {'type': 'content_block_stop', 'index': 0})
            ev('message_delta', {'type': 'message_delta', 'delta': {'stop_reason': msg['stop_reason'], 'stop_sequence': None}, 'usage': {'output_tokens': msg['usage']['output_tokens']}})
            ev('message_stop', {'type': 'message_stop'})
            return
        payload = json.dumps(msg).encode()
        self.send_response(200)
        self.send_header('content-type', 'application/json')
        self.send_header('request-id', 'req_stub')
        self.end_headers()
        self.wfile.write(payload)


if __name__ == '__main__':
    print(f'LLM stub on http://127.0.0.1:{PORT}', flush=True)
    ThreadingHTTPServer(('127.0.0.1', PORT), Handler).serve_forever()