"""Local Windows Server collector. No remoting, credentials or shell interpolation."""
import argparse
import ctypes
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import platform
import struct
import subprocess
import time

from jsonschema import Draft202012Validator, FormatChecker
from samsara.core.process import CollectionCancelled, ProcessFailure, capture
from samsara.core.privacy import minimize_inventory, load_privacy_policy, safe_time
from samsara.core.detection import check_profile
from samsara.models.inventory import inventory_schema, validate_inventory

SCRIPT = Path(__file__).resolve().parents[2] / 'powershell' / 'Get-ServerInventory.ps1'
SECTIONS = ('system', 'roles', 'services', 'network', 'tasks', 'software', 'storage', 'shares', 'certificates')
SERVER_BUILDS = {'14393': '2016', '17763': '2019', '20348': '2022', '26100': '2025'}
MAX_OUTPUT = 10 * 1024 * 1024


class WindowsCollectionError(RuntimeError):
    def __init__(self, code, message):
        super().__init__(message)
        self.code = code


def timestamp():
    return datetime.now(timezone.utc).isoformat()


def powershell_executable():
    if platform.system() != 'Windows':
        raise ValueError('Local Windows collection requires an authorized Windows Server execution host.')
    if struct.calcsize('P') != 8:
        raise ValueError('Use 64-bit Python for Windows Server collection.')
    # Query Windows itself; never look up powershell.exe in PATH or trust SystemRoot.
    buffer = ctypes.create_unicode_buffer(32768)
    size = ctypes.windll.kernel32.GetSystemDirectoryW(buffer, len(buffer))
    if size == 0 or size >= len(buffer):
        raise ValueError('Unable to locate the Windows system directory.')
    executable = Path(buffer.value) / 'WindowsPowerShell' / 'v1.0' / 'powershell.exe'
    if not executable.is_file() or not SCRIPT.is_file():
        raise ValueError('Windows PowerShell 5.1 or bundled collector scripts are unavailable.')
    return executable


def run_section(section, executable, cancelled=lambda: False, timeout=45, diagnostic=False):
    if section not in SECTIONS:
        raise ValueError('Unknown Windows inventory section.')
    if cancelled():
        raise CollectionCancelled()
    args = [str(executable), '-NoLogo', '-NoProfile', '-NonInteractive', '-File', str(SCRIPT), '-Section', section]
    if diagnostic:
        args.append('-Diagnostic')
    environment = {key: os.environ[key] for key in ('SystemRoot','WINDIR','TEMP','TMP','COMSPEC','SystemDrive') if key in os.environ}
    environment['PSModulePath'] = str(Path(executable).parent / 'Modules')
    try:
        raw = capture(args, env=environment, timeout=timeout, cancelled=cancelled,
                      limit=MAX_OUTPUT, creationflags=getattr(subprocess,'CREATE_NO_WINDOW',0))
    except ProcessFailure as exc:
        if exc.returncode == 3:
            raise WindowsCollectionError('serialization_failed', 'PowerShell could not serialize inventory metadata; no successful evidence was substituted.') from None
        raise WindowsCollectionError(exc.code, str(exc)) from None
    try:
        return json.loads(raw.decode('utf-8-sig'))
    except (UnicodeError, ValueError, RecursionError):
        raise WindowsCollectionError('invalid_output','PowerShell did not return a valid JSON section.') from None


def validate_section(record, section):
    schema = inventory_schema()
    properties = dict(schema['properties']['checks']['additionalProperties']['properties'])
    properties['section'] = {'const': section}
    if section == 'system':
        properties['data'] = {'type': 'object', 'required': ['hostname','platform','os_build','product_type','os_version'],
                              'properties': schema['properties']['server']['properties']}
    else:
        properties['data'] = schema['properties']['inventory']['properties']['storage' if section == 'shares' else section]
    section_schema = {'type': 'object', 'required': ['section','status','message','started_at','completed_at','errors','permission_limitations','limitations','data'], 'properties': properties}
    if not Draft202012Validator(section_schema, format_checker=FormatChecker()).is_valid(record):
        raise WindowsCollectionError('invalid_output', 'PowerShell returned a section that violates the inventory contract.')
    start, end = safe_time(record['started_at']), safe_time(record['completed_at'])
    if start is None or end is None or end < start:
        # Both native timestamps use the same UTC format. This is not a trust check.
        raise WindowsCollectionError('invalid_output', 'PowerShell returned inconsistent section timestamps.')
    return record


def verify_server(server, profile='auto'):
    year = SERVER_BUILDS.get(server.get('os_build'))
    if server.get('platform') != 'windows' or server.get('product_type') not in (2, 3) or year is None:
        raise ValueError('Unable to verify Windows Server 2016, 2019, 2022 or 2025. Workstation editions and unknown builds are not supported.')
    if server.get('os_version') != year:
        raise ValueError('Windows Server identity fields are inconsistent; no assessment was accepted.')
    check_profile({'family': 'windows', 'version': year, 'name': server.get('operating_system', 'Windows Server ' + year)}, profile)


def failure_check(code, message, section, start, status='failed'):
    return {'status': status, 'message': message, 'started_at': start, 'completed_at': timestamp(),
            'errors': [{'code': code, 'scope': section, 'message': message}],
            'permission_limitations': ['Permissions could not be verified for this check.'],
            'limitations': ['No successful evidence was substituted for this failure.']}


def collect_local_windows(stage=lambda message: None, cancelled=lambda: False, profile='auto',
                          section_timeout=45, total_timeout=300, runner=None, diagnostic=False):
    if platform.system() != 'Windows':
        raise ValueError('Local Windows collection requires an authorized Windows Server execution host.')
    if profile != 'auto' and not profile.startswith('windows:'):
        raise ValueError('Select Auto-detect or a Windows Server profile for this execution host.')
    load_privacy_policy()
    executable = powershell_executable()
    runner = runner or run_section
    started = time.monotonic()
    result = {'schema_version': '1.0', 'synthetic': False, 'collected_at': timestamp(), 'server': {},
              'checks': {}, 'inventory': {key: [] for key in ('roles','services','network','tasks','software','storage','certificates')}}
    total_bytes = 0
    for section in SECTIONS:
        if cancelled():
            raise CollectionCancelled()
        stage('Collecting local Windows ' + section + '…')
        start = timestamp()
        remaining = total_timeout - (time.monotonic() - started)
        if remaining <= 0:
            result['checks'][section] = failure_check('overall_timeout', 'Overall collection time limit reached before this check.', section, start, 'not_collected')
            if section == 'system':
                raise ValueError('The Windows identity check could not run within the time limit.')
            continue
        try:
            options = {'cancelled': cancelled, 'timeout': min(section_timeout, remaining)}
            if diagnostic:
                options['diagnostic'] = True
            record = validate_section(runner(section, executable, **options), section)
            total_bytes += len(json.dumps(record).encode('utf-8'))
            if total_bytes > MAX_OUTPUT:
                raise WindowsCollectionError('output_limit', 'Total inventory exceeded the 10 MiB limit; this section was omitted.')
            check = {key: record[key] for key in ('status','message','started_at','completed_at','errors','permission_limitations','limitations')}
            result['checks'][section] = check
            if section == 'system':
                if record['status'] == 'failed':
                    raise ValueError('Windows identity collection failed. Verify local CIM read access and Windows PowerShell policy.')
                verify_server(record['data'], profile)
                result['server'] = record['data']
                result['server']['collection_source'] = 'Local Windows read-only collection'
            elif section == 'shares':
                result['inventory']['storage'].extend(record['data'])
            else:
                result['inventory'][section] = record['data']
        except WindowsCollectionError as exc:
            result['checks'][section] = failure_check(exc.code, str(exc), section, start)
            if section == 'system':
                raise ValueError('Windows identity collection failed: ' + str(exc)) from None
    if cancelled():
        raise CollectionCancelled()
    return validate_inventory(minimize_inventory(result, trusted_source='Local Windows read-only collection'))


def main():
    parser = argparse.ArgumentParser(description='Read-only local Windows Server inventory. No remoting is used.')
    parser.add_argument('--output', required=True, help='New local JSON file to create')
    parser.add_argument('--diagnostic', action='store_true', help='Include sanitized exception types and categories')
    parser.add_argument('--profile', choices=['auto','windows:2016','windows:2019','windows:2022','windows:2025'], default='auto')
    args = parser.parse_args()
    try:
        data = collect_local_windows(profile=args.profile, diagnostic=args.diagnostic)
        from samsara.core.storage import write_private
        write_private(args.output, json.dumps(data, indent=2))
        print('Windows inventory saved. Review all failed/partial checks and permission limitations before retirement decisions.')
    except (OSError, ValueError, WindowsCollectionError):
        parser.exit(1, 'Windows inventory was not saved. Check supported local Server OS, Python architecture, script policy, permissions and a new writable output filename.\n')


if __name__ == '__main__':
    main()