#!/usr/bin/env python3
"""Demostración limitada de herramientas para Ollama, Python 3 sin dependencias.

En el pod por SSH: python3 /workspace/demo_tools.py
Prueba local sin contactar IA: python3 demo_tools.py --self-check
Requiere ffmpeg y ffprobe. API privada local: http://127.0.0.1:11434.
El modelo solo puede pedir un clip fijo; nunca se ejecutan comandos del modelo.
"""
import argparse
import json
import pathlib
import shutil
import subprocess
import tempfile
import urllib.request

MODEL = 'qwen3.5:9b'
ENDPOINT = 'http://127.0.0.1:11434/api/chat'
EXPECTED = {'duration_seconds': 5, 'width': 1280, 'height': 720}
TOOL = {'type': 'function', 'function': {
    'name': 'render_video',
    'description': 'Crea un clip MP4 con FFmpeg. Devuelve ruta, duración y resolución verificadas.',
    'parameters': {'type': 'object', 'properties': {
        k: {'type': 'integer', 'enum': [v]} for k, v in EXPECTED.items()
    }, 'required': list(EXPECTED), 'additionalProperties': False}
}}


def validate(arguments):
    if not isinstance(arguments, dict) or set(arguments) != set(EXPECTED):
        raise ValueError('Se requieren exactamente duration_seconds, width y height.')
    if any(type(arguments[k]) is not int or arguments[k] != v for k, v in EXPECTED.items()):
        raise ValueError('Solo se permite un clip de 5 segundos y 1280 × 720.')


def render_video(arguments, output):
    validate(arguments)
    for binary in ('ffmpeg', 'ffprobe'):
        if not shutil.which(binary):
            raise RuntimeError(f'Falta {binary}.')
    # Todos los argumentos ejecutables son constantes; el modelo no elige rutas ni filtros.
    subprocess.run([
        'ffmpeg', '-hide_banner', '-loglevel', 'error', '-y',
        '-f', 'lavfi', '-i', 'color=c=black:s=1280x720:r=30:d=5',
        '-f', 'lavfi', '-i', 'color=c=white:s=180x180:r=30:d=5',
        '-filter_complex', "[0:v][1:v]overlay=x='100+180*t':y=270:shortest=1,format=yuv420p",
        '-t', '5', '-c:v', 'libx264', '-preset', 'veryfast', str(output)
    ], check=True, timeout=120, capture_output=True)
    probe = json.loads(subprocess.check_output([
        'ffprobe', '-v', 'error', '-show_entries',
        'format=duration:stream=width,height,codec_name', '-of', 'json', str(output)
    ], text=True, timeout=30))
    streams = probe.get('streams', [])
    duration = float(probe['format']['duration'])
    if (len(streams) != 1 or streams[0].get('width') != 1280
            or streams[0].get('height') != 720 or streams[0].get('codec_name') != 'h264'
            or abs(duration - 5) > .05 or output.stat().st_size == 0):
        raise RuntimeError('El archivo no pasó la verificación.')
    return {'path': str(output), 'duration_seconds': duration,
            'width': 1280, 'height': 720, 'codec': 'h264', 'verified': True}


def chat(messages, tools=None):
    data = {'model': MODEL, 'messages': messages, 'stream': False, 'think': False,
            'keep_alive': '15m',
            'options': {'num_ctx': 65536, 'num_predict': 600, 'temperature': 0}}
    if tools is not None:
        data['tools'] = tools
    request = urllib.request.Request(ENDPOINT, data=json.dumps(data).encode(),
                                     headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(request, timeout=240) as response:
        result = json.load(response)
    if 'error' in result:
        raise RuntimeError(result['error'])
    return result['message']


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--self-check', action='store_true')
    args = parser.parse_args()
    if args.self_check:
        # Comprueba límites y FFmpeg; no hace solicitudes HTTP.
        for invalid in ({}, dict(EXPECTED, width=1920), dict(EXPECTED, width=True),
                        dict(EXPECTED, command='echo')):
            try:
                validate(invalid)
            except ValueError:
                pass
            else:
                raise RuntimeError('Falló la validación de límites.')
        with tempfile.TemporaryDirectory(prefix='ollama-demo-') as scratch:
            result = render_video(EXPECTED, pathlib.Path(scratch) / 'demo.mp4')
            print(json.dumps({k: v for k, v in result.items() if k != 'path'}, ensure_ascii=False))
        print('Self-check correcto. Archivo temporal eliminado. No se contactó a la IA.')
        return
    output = pathlib.Path('/workspace/tutorial-demo.mp4')
    if not output.parent.is_dir():
        raise RuntimeError('Ejecuta este comando dentro del pod, donde existe /workspace.')
    messages = [{'role': 'user', 'content':
        'Usa la herramienta render_video para crear un video de 5 segundos, 1280 por 720. '
        'No afirmes que existe hasta recibir el resultado de la herramienta.'}]
    assistant = chat(messages, [TOOL])
    calls = assistant.get('tool_calls', [])
    if len(calls) != 1 or calls[0].get('function', {}).get('name') != 'render_video':
        raise RuntimeError('El modelo no solicitó la única herramienta permitida. No se ejecutó nada.')
    arguments = calls[0]['function'].get('arguments')
    validate(arguments)
    print('El modelo solicitó render_video:', json.dumps(arguments))
    result = render_video(arguments, output)
    print('Resultado verificado:', json.dumps(result, ensure_ascii=False))
    messages += [assistant, {'role': 'tool', 'tool_name': 'render_video',
                            'content': json.dumps(result, ensure_ascii=False)}]
    final = chat(messages)
    print('Respuesta del modelo:', final.get('content', ''))


if __name__ == '__main__':
    try:
        main()
    except (OSError, ValueError, RuntimeError, KeyError, subprocess.SubprocessError) as error:
        raise SystemExit(f'No se completó la demostración: {error}')