#!/usr/bin/env python3
"""Verify local code, tests, databases and Docker without calling an external LLM."""
from __future__ import annotations

import argparse
import ast
import asyncio
import ipaddress
import json
import os
from pathlib import Path
import re
import shlex
import socket
import sqlite3
import subprocess
import sys
import time
import tempfile
import types
import urllib.request
import zipfile

ROOT = Path(__file__).resolve().parents[1]


def local_test_worker() -> int:
    """Allow live local API/Neo4j tests, forbid outbound model calls."""
    os.environ.pop('OPENWEBUI_API_KEY', None)
    for line in (ROOT / '.env').read_text().splitlines():
        key, separator, value = line.partition('=')
        if separator and key in {'NEO4J_URI', 'NEO4J_USER', 'NEO4J_PASSWORD', 'NEO4J_DATABASE'}:
            values = shlex.split(value)
            if values:
                os.environ[key] = values[0]
    for key in ['RUN_NEO4J_TESTS', 'RUN_LIVE_TESTS', 'RUN_LIVE_RAG_E2E', 'HF_HUB_OFFLINE', 'TRANSFORMERS_OFFLINE']:
        os.environ[key] = '1'
    original_connect, original_connect_ex = socket.socket.connect, socket.socket.connect_ex
    original_dns = socket.getaddrinfo

    def local_ports():
        # HTTP transport tests register only their ephemeral loopback listener.
        extra = os.getenv('ZP_TEST_HTTP_PORT')
        return {7687, 8000} | ({int(extra)} if extra else set())

    def permitted(address):
        try:
            return address[1] in local_ports() and ipaddress.ip_address(address[0]).is_loopback
        except (ValueError, TypeError, IndexError):
            return False

    def guarded_connect(self, address):
        if self.family in {socket.AF_INET, socket.AF_INET6} and not permitted(address):
            raise RuntimeError('External network disabled during local verification')
        return original_connect(self, address)

    def guarded_connect_ex(self, address):
        if self.family in {socket.AF_INET, socket.AF_INET6} and not permitted(address):
            raise RuntimeError('External network disabled during local verification')
        return original_connect_ex(self, address)

    def guarded_dns(host, port, *args, **kwargs):
        if host not in {'localhost', '127.0.0.1', '::1'} or int(port) not in local_ports():
            raise RuntimeError('External DNS disabled during local verification')
        return original_dns(host, port, *args, **kwargs)

    socket.socket.connect, socket.socket.connect_ex = guarded_connect, guarded_connect_ex
    socket.getaddrinfo = guarded_dns
    import pytest
    return pytest.main(['-q', '-p', 'no:cacheprovider', 'test', 'tools/rag/tests', 'tools/graphrag/tests'])


def verify_code() -> dict:
    files = [path for name in ['app', 'tools', 'packages', 'test', 'scripts'] for path in (ROOT / name).rglob('*.py')]
    for path in files:
        text = path.read_text()
        compile(text, str(path), 'exec')
        for node in ast.walk(ast.parse(text)):
            modules = ([node.module] if isinstance(node, ast.ImportFrom) and node.module else
                       [item.name for item in node.names] if isinstance(node, ast.Import) else [])
            for module in modules:
                if module.split('.')[0] in {'tools', 'app'}:
                    target = ROOT.joinpath(*module.split('.'))
                    if not target.with_suffix('.py').is_file() and not (target / '__init__.py').is_file():
                        raise RuntimeError(f'Unresolved internal module {module} in {path.relative_to(ROOT)}')
    shells = [path for name in ['scripts', 'tools'] for path in (ROOT / name).rglob('*.sh')]
    for path in shells:
        subprocess.run(['bash', '-n', str(path)], check=True)
    return {'compiled_python_files': len(files), 'checked_shell_scripts': len(shells)}


def verify_database() -> dict:
    path = ROOT / 'data/zp_index.sqlite'
    with sqlite3.connect(path.resolve().as_uri() + '?mode=ro', uri=True) as conn:
        assert conn.execute('PRAGMA quick_check').fetchone()[0] == 'ok', 'SQLite integrity failed'
        assert not conn.execute('PRAGMA foreign_key_check').fetchall(), 'SQLite foreign keys failed'
        counts = {table: conn.execute(f'SELECT count(*) FROM {table}').fetchone()[0]
                  for table in ['documents', 'chunks', 'chunks_fts', 'chunk_embeddings']}
        assert counts['chunks'] == counts['chunks_fts'] and counts['chunks'] > 0, 'FTS count mismatch'
    return counts


def verify_chunk_limits() -> dict:
    import numpy as np
    from tools.rag.build_chunks import TokenCounter
    from tools.rag.embedding_utils import build_embedding_text, split_passage, passage_token_count, get_embedding_model, embedding_model_name
    chunks = json.loads((ROOT/'data/chunks.json').read_text())
    if isinstance(chunks, dict):
        chunks = chunks['chunks']
    counter = TokenCounter(os.getenv('CHUNK_TOKEN_ENCODING', 'cl100k_base'))
    limit = int(os.getenv('CHUNK_MAX_TOKENS', '450'))
    model = get_embedding_model(embedding_model_name())
    maximum = 0
    expected_segments = set()
    with sqlite3.connect((ROOT/'data/zp_index.sqlite').resolve().as_uri()+'?mode=ro',uri=True) as conn:
        stored = dict(conn.execute('SELECT chunk_id,text FROM chunks'))
        assert len(stored) == len(chunks)
        for chunk in chunks:
            assert stored[chunk['chunk_id']] == chunk['text']
            size = counter.count(chunk['text'])
            maximum = max(maximum, size)
            assert size <= limit, 'Oversized source chunk'
            original = build_embedding_text(chunk)
            parts = split_passage(original)
            assert ''.join(parts) == original, 'Embedding segmentation loses content'
            for index, part in enumerate(parts):
                assert passage_token_count(part) <= model.max_seq_length, 'Oversized embedding input'
                expected_segments.add((chunk['chunk_id'],index))
        actual_segments = set()
        for cid,index,dimensions,blob in conn.execute('SELECT chunk_id,segment_index,dimensions,embedding FROM chunk_embeddings'):
            vector = np.frombuffer(blob,dtype=np.float32)
            assert len(vector)==dimensions and np.isfinite(vector).all()
            assert np.isclose(np.linalg.norm(vector),1.0,atol=0.001), 'Invalid embedding vector'
            actual_segments.add((cid,index))
        assert actual_segments == expected_segments, 'Embedding coverage mismatch'
    return {'chunks':len(chunks),'max_chunk_tokens':maximum,'chunk_limit':limit,
            'embedding_segments':len(actual_segments),'embedding_limit':model.max_seq_length,
            'source_content_preserved':True}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--docker-context', default='default')
    parser.add_argument('--git-network', action='store_true', help='Read-only probe of wiki origin over its configured SSH connection')
    parser.add_argument('--wait-ready-seconds', type=int, default=90,
                        help='Bounded wait for API/model warm-up after Docker startup')
    parser.add_argument('--output', type=Path)
    parser.add_argument('--no-save', action='store_true', help='Check in memory without writing reports or logs')
    parser.add_argument('--test-worker', action='store_true', help=argparse.SUPPRESS)
    args = parser.parse_args()
    os.chdir(ROOT)
    sys.path.insert(0, str(ROOT))
    if args.test_worker:
        return local_test_worker()
    from tools.runtime.report_files import report_output, report_log, retention_plan, apply_moves
    if args.no_save and args.output:
        parser.error('--no-save cannot be combined with --output')
    args.output = None if args.no_save else report_output(args.output, 'verification')
    report = {'checks': {}, 'external_llm_calls': False}
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)

    def check(name, action):
        started = time.monotonic()
        try:
            result = action()
            record = {'passed': True, 'details': result}
        except Exception as error:
            record = {'passed': False, 'error_type': type(error).__name__}
        record['seconds'] = round(time.monotonic() - started, 2)
        report['checks'][name] = record
        if args.output:
            args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + '\n')
        print(f'{name}: {"OK" if record["passed"] else "FAILED"}', flush=True)

    def command(*argv):
        result = subprocess.run(argv, text=True, capture_output=True, timeout=90)
        if result.returncode:
            raise RuntimeError('Verification command failed; inspect the service locally')
        return result.stdout.strip()

    docker = ['docker', '--context', args.docker_context]
    compose = [*docker, 'compose', '--project-name', 'zp-agent',
               '--project-directory', str(ROOT / 'deployment'),
               '--env-file', str(ROOT / '.env'),
               '-f', str(ROOT / 'deployment/docker-compose.yml')]
    check('python_and_shell', verify_code)
    check('git_diff', lambda: command('git', 'diff', '--check'))
    check('git_integrity', lambda: command('git', 'fsck', '--connectivity-only'))
    check('compose_configuration', lambda: command(*compose, 'config', '--quiet'))

    def health():
        deadline = time.monotonic() + max(0, min(args.wait_ready_seconds, 300))
        while True:
            try:
                with urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3) as response:
                    value = json.load(response)
                assert value['ready'] and all(tool['ready'] for tool in value['tools'].values())
                break
            except (OSError, AssertionError):
                if time.monotonic() >= deadline:
                    raise
                time.sleep(1)
        assert value['tools']['rag']['index_version'] == value['tools']['graphrag']['index_version']
        return {'ready': True, 'index_version': value['index_version'],
                'tools': {name: item['ready'] for name, item in value['tools'].items()}}

    check('api_and_both_indexes', health)
    check('sqlite_integrity', verify_database)
    check('chunk_and_embedding_limits', verify_chunk_limits)
    check('container_readiness', lambda: command(*compose, 'exec', '-T', 'zp-agent-api',
        'python', '-m', 'tools.runtime.runtime_health'))
    if args.git_network:
        probe = ("import os,subprocess; from tools.sync.wiki_git import expected_branch; "
                 "p=subprocess.run(['git','-C',os.environ['ZPWIKI_ROOT'],'ls-remote','--exit-code',"
                 "'origin','refs/heads/'+expected_branch()],capture_output=True,text=True,timeout=60); "
                 "assert p.returncode==0,'Wiki origin unavailable'; print(p.stdout.strip())")
        check('wiki_git_ssh', lambda: command(*compose, 'exec', '-T', 'zp-agent-api', 'python', '-c', probe))

    def tool_packages():
        from scripts.build_tool_packages import build_package
        key = os.getenv('SEARCH_API_KEY', '').strip()
        if not key:
            for line in (ROOT / '.env').read_text().splitlines():
                name, separator, value = line.partition('=')
                if separator and name == 'SEARCH_API_KEY':
                    key = shlex.split(value)[0]
        assert key, 'Search API key is missing'
        built = []

        async def run(package_dir):
            graph_tool = None
            versions = []
            for engine in ('rag', 'graphrag'):
                package = build_package(engine, package_dir)
                with zipfile.ZipFile(package['archive']) as archive:
                    form = json.loads(archive.read('openwebui-tool.json'))[0]
                module = types.ModuleType('verify_packaged_' + engine)
                sys.modules[module.__name__] = module
                try:
                    exec(compile(form['content'], '<packaged-tool>', 'exec'), module.__dict__)
                    tool = module.Tools()
                    tool.valves.BACKEND_URL = 'http://127.0.0.1:8000'
                    tool.valves.BACKEND_API_KEY = key
                    method = getattr(tool, 'retrieve_zpwiki_context' if engine == 'rag' else 'retrieve_zpwiki_graph')
                    result = json.loads(await method('Aký je názov diplomovej práce osoby Ján Holp?'))
                    assert result['status'] == 'ok' and result['sources'], 'Packaged tool cannot retrieve evidence'
                    assert result['requires_final_answer_verification'] is True
                    versions.append(result['index_version'])
                    if engine == 'graphrag':
                        graph_tool = tool
                    built.append({name: package[name] for name in ['package_id', 'archive', 'sha256']})
                finally:
                    sys.modules.pop(module.__name__, None)
            assert versions[0] == versions[1], 'Tools use different data versions'
            negative = json.loads(await graph_tool.retrieve_zpwiki_graph(
                'Aký je názov diplomovej práce osoby Ján Holp v roku 1999?'))
            assert negative['status'] == 'no_answer' and not negative['facts']
            graph_tool.valves.BACKEND_API_KEY = 'invalid-test-key'
            assert json.loads(await graph_tool.retrieve_zpwiki_graph('Práca?'))['reason'] == 'unauthorized'
        if args.no_save:
            with tempfile.TemporaryDirectory(prefix='zp-verify-packages-') as folder:
                asyncio.run(run(Path(folder)))
        else:
            asyncio.run(run(ROOT / 'data/tool_packages'))
        return {'packages': built, 'live_retrieval': True, 'same_data_version': True,
                'negative_query': True, 'authorization': True, 'external_llm_calls': False}

    check('openwebui_tool_packages', tool_packages)

    def tests():
        environment = {**os.environ}
        environment.pop('OPENWEBUI_API_KEY', None)
        environment['PYTHONDONTWRITEBYTECODE'] = '1'
        command = [sys.executable, str(Path(__file__).resolve()), '--test-worker']
        if args.no_save:
            result = subprocess.run(command, env=environment, capture_output=True, text=True, timeout=900)
            text = result.stdout + result.stderr
            if result.returncode:
                print(text[-16000:], flush=True)
            log = None
        else:
            log = report_log(args.output)
            with log.open('w') as output:
                result = subprocess.run(command, env=environment, stdout=output, stderr=subprocess.STDOUT, timeout=900)
            text = log.read_text()
        assert result.returncode == 0, 'Regression failed; see console or selected pytest log'
        counts = {name: int(value) for value, name in re.findall(r'(\d+) (passed|failed|skipped)', text)}
        assert counts.get('passed', 0) > 0 and counts.get('skipped', 0) == 0
        print('Regression: ' + str(counts), flush=True)
        return {**counts, 'log': str(log) if log else None, 'external_network_blocked': True}

    check('full_local_regression', tests)
    report['passed'] = all(item['passed'] for item in report['checks'].values())
    if args.output:
        args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + '\n')
        moves = retention_plan()
        if moves:
            apply_moves(moves)
        print('Report: ' + str(args.output))
    return 0 if report['passed'] else 1


if __name__ == '__main__':
    raise SystemExit(main())