"""Autarch Android Forensics AI — threat analysis, CVE lookup, heuristic scanning, and alert receiver."""

from datetime import datetime
from flask import Blueprint, request, jsonify
from web.auth import login_required

android_forensics_ai_bp = Blueprint('android_forensics_ai', __name__, url_prefix='/android-forensics/ai')

_alerts = []   # in-memory alert store (last 500)
_MAX_ALERTS = 500


def _mgr():
    from core.android_forensics import get_android_forensics_manager
    return get_android_forensics_manager()


def _cve():
    from core.cve_db import get_cve_db
    return get_cve_db()


def _llm():
    from core.llm import get_llm
    return get_llm()


def _body():
    return request.get_json(silent=True) or {}


# ── AI Analysis ──────────────────────────────────────────────────────

@android_forensics_ai_bp.route('/analyze', methods=['POST'])
@login_required
def analyze():
    data = _body()
    serial   = data.get('serial', '').strip()
    acq_id   = data.get('acq_id', '').strip()
    backend  = data.get('backend', 'autarch')
    focus    = data.get('focus', ['all'])

    try:
        # Load acquisition report from disk, fall back to live acquisition
        report = None
        if serial and acq_id:
            exported = _mgr().export_report(serial, acq_id)
            report = exported.get('report') if exported else None

        if not report:
            report = _mgr().full_acquisition(serial) if serial else {}

        # Build context dict from report
        packages  = report.get('packages', [])
        processes = report.get('processes', {})
        proc_output = processes.get('output', '') if isinstance(processes, dict) else str(processes)
        proc_lines  = proc_output.splitlines()[:50]

        props       = report.get('getprop', {})
        model           = props.get('ro.product.model', 'Unknown') if isinstance(props, dict) else 'Unknown'
        android_version = props.get('ro.build.version.release', 'Unknown') if isinstance(props, dict) else 'Unknown'

        network_info = report.get('network', {})
        ioc_findings = report.get('ioc_findings', report.get('findings', []))

        pkg_list   = packages[:100] if isinstance(packages, list) else []
        user_pkgs  = [p for p in pkg_list if not p.get('system', False)] if pkg_list and isinstance(pkg_list[0], dict) else pkg_list
        pkg_names  = [p.get('package', str(p)) if isinstance(p, dict) else str(p) for p in pkg_list]

        context_dict = {
            'model':           model,
            'android_version': android_version,
            'pkg_count':       len(pkg_list),
            'user_count':      len(user_pkgs),
            'pkg_sample':      ', '.join(pkg_names[:20]),
            'proc_sample':     '\n'.join(proc_lines),
            'network_info':    str(network_info)[:500],
            'ioc_findings':    str(ioc_findings)[:500],
        }

        if backend == 'raw':
            return jsonify({'ok': True, 'mode': 'raw', 'dump': context_dict})

        if backend == 'ondevice':
            prompt_str = (
                f"Device: {context_dict['model']} Android {context_dict['android_version']}. "
                f"Packages ({context_dict['pkg_count']}): {context_dict['pkg_sample']}. "
                f"Processes: {context_dict['proc_sample'][:200]}. "
                f"Network: {context_dict['network_info'][:200]}. "
                f"IOC findings: {context_dict['ioc_findings'][:200]}. "
                "Identify threats and return risk assessment."
            )
            return jsonify({
                'ok':     True,
                'mode':   'ondevice',
                'prompt': prompt_str,
                'schema': {
                    'risk_level': 'string',
                    'findings':   'array',
                    'summary':    'string',
                },
            })

        # autarch backend — full LLM analysis
        prompt = (
            "You are a mobile security analyst. Analyze this Android forensic acquisition and identify threats.\n\n"
            f"Device info: {context_dict['model']} Android {context_dict['android_version']}\n"
            f"Installed packages ({context_dict['pkg_count']} total, {context_dict['user_count']} user apps): {context_dict['pkg_sample']}\n"
            f"Running processes (sample): {context_dict['proc_sample']}\n"
            f"Network properties: {context_dict['network_info']}\n"
            f"IOC pre-scan findings: {context_dict['ioc_findings']}\n\n"
            'Return JSON: {"risk_level": "critical|high|medium|low|clean", '
            '"findings": [{"severity": "...", "type": "...", "detail": "...", "recommendation": "..."}], '
            '"summary": "one paragraph"}'
        )

        response = _llm().generate(prompt, max_tokens=1000)
        try:
            import json
            parsed_result = json.loads(response)
        except Exception:
            parsed_result = response

        return jsonify({'ok': True, 'mode': 'autarch', 'analysis': parsed_result})

    except Exception as e:
        return jsonify({'ok': False, 'error': str(e)})


# ── CVE Check ────────────────────────────────────────────────────────

@android_forensics_ai_bp.route('/cve-check', methods=['POST'])
@login_required
def cve_check():
    data        = _body()
    packages    = data.get('packages', [])
    sdk_version = data.get('sdk_version')

    try:
        _severity_order = {'critical': 0, 'high': 1, 'medium': 2, 'low': 3, 'unknown': 4}

        findings = []

        pkg_results = _cve().check_packages_batch(packages)
        for pkg_name, cves in (pkg_results or {}).items():
            for cve in (cves or []):
                findings.append({
                    'package':         pkg_name,
                    'cve_id':          cve.get('cve_id', ''),
                    'severity':        cve.get('severity', 'unknown'),
                    'description':     cve.get('description', ''),
                    'patched_version': cve.get('patched_version', ''),
                    'source':          cve.get('source', ''),
                })

        if sdk_version:
            sdk_cves = _cve().get_cve_for_android_sdk(sdk_version) or []
            for cve in sdk_cves:
                findings.append({
                    'package':         f'android-sdk:{sdk_version}',
                    'cve_id':          cve.get('cve_id', ''),
                    'severity':        cve.get('severity', 'unknown'),
                    'description':     cve.get('description', ''),
                    'patched_version': cve.get('patched_version', ''),
                    'source':          cve.get('source', ''),
                })

        findings.sort(key=lambda f: _severity_order.get(f['severity'].lower(), 4))

        return jsonify({
            'ok':               True,
            'findings':         findings,
            'total':            len(findings),
            'packages_checked': len(packages),
        })

    except Exception as e:
        return jsonify({'ok': False, 'error': str(e)})


# ── Classify ─────────────────────────────────────────────────────────

@android_forensics_ai_bp.route('/classify', methods=['POST'])
@login_required
def classify():
    data       = _body()
    raw_output = data.get('raw_output', '')
    data_type  = data.get('type', 'packages')

    try:
        _type_context = {
            'packages':  'a list of installed Android packages',
            'logcat':    'Android logcat output',
            'processes': 'a list of running Android processes',
            'getprop':   'Android system properties (getprop output)',
        }
        context_label = _type_context.get(data_type, f'Android {data_type} data')

        prompt = (
            f"You are a mobile security analyst. Classify the following {context_label} "
            "for security threats, stalkerware, spyware, adware, or other malicious indicators. "
            "Provide a brief threat classification summary.\n\n"
            f"Data:\n{raw_output[:3000]}"
        )

        response = _llm().generate(prompt, max_tokens=500)
        return jsonify({'ok': True, 'classification': response})

    except Exception as e:
        return jsonify({'ok': False, 'error': str(e)})


# ── Heuristic Scan ───────────────────────────────────────────────────

@android_forensics_ai_bp.route('/heuristic', methods=['POST'])
@login_required
def heuristic():
    data       = _body()
    data_type  = data.get('data_type', 'packages')
    items      = data.get('items', [])
    threshold  = float(data.get('threshold', 0.6))

    _prompts = {
        'packages': (
            'Classify each Android package name for stalkerware/spyware/adware risk. '
            'Score 0.0-1.0. High risk: names impersonating system apps, random strings, MDM-sounding names. '
            'Return JSON array: [{"item": "...", "score": 0.0, "reason": "...", "category": "stalkerware|adware|suspicious|clean"}]'
        ),
        'processes': (
            'Classify each process name for malware risk. Score 0.0-1.0. '
            'Suspicious: random strings, names mimicking system processes. '
            'Return JSON array: [{"item": "...", "score": 0.0, "reason": "...", "category": "malware|suspicious|clean"}]'
        ),
        'network': (
            'Classify each domain/IP for C2/malware beacon risk. Score 0.0-1.0. '
            'High risk: DGA patterns, dynamic DNS abuse, unusual TLDs. '
            'Return JSON array: [{"item": "...", "score": 0.0, "reason": "...", "category": "c2_beacon|dga_domain|suspicious|clean"}]'
        ),
        'files': (
            'Classify each file path for spyware/stalkerware artifact risk. Score 0.0-1.0. '
            'Return JSON array: [{"item": "...", "score": 0.0, "reason": "...", "category": "spyware|suspicious|clean"}]'
        ),
    }

    try:
        llm = _llm()
        if llm is None:
            return jsonify({'ok': False, 'error': 'LLM not loaded', 'findings': []})

        base_prompt = _prompts.get(data_type, _prompts['packages'])
        chunk_size  = 50
        all_results = []

        for i in range(0, len(items), chunk_size):
            chunk = items[i:i + chunk_size]
            chunk_text = '\n'.join(str(it) for it in chunk)
            prompt = f"{base_prompt}\n\nItems:\n{chunk_text}"

            try:
                response = llm.generate(prompt, max_tokens=1000)
                import json
                parsed = json.loads(response)
                if isinstance(parsed, list):
                    all_results.extend(parsed)
            except Exception:
                # Skip chunks that fail to parse
                continue

        filtered = [r for r in all_results if isinstance(r, dict) and float(r.get('score', 0)) >= threshold]

        return jsonify({
            'ok':           True,
            'findings':     filtered,
            'total_scanned': len(items),
            'flagged':      len(filtered),
        })

    except Exception as e:
        return jsonify({'ok': False, 'error': str(e), 'findings': []})


# ── Alert Receiver ───────────────────────────────────────────────────

@android_forensics_ai_bp.route('/alert', methods=['POST'])
@login_required
def alert():
    data = _body()
    entry = {
        'serial':     data.get('serial', ''),
        'event_type': data.get('event_type', ''),
        'severity':   data.get('severity', ''),
        'detail':     data.get('detail', ''),
        'timestamp':  data.get('timestamp', datetime.utcnow().isoformat()),
    }
    _alerts.append(entry)
    if len(_alerts) > _MAX_ALERTS:
        del _alerts[:-_MAX_ALERTS]
    return jsonify({'ok': True, 'alert_id': len(_alerts)})


# ── Alerts List ──────────────────────────────────────────────────────

@android_forensics_ai_bp.route('/alerts/list', methods=['GET'])
@login_required
def alerts_list():
    return jsonify({'ok': True, 'alerts': _alerts[-100:], 'total': len(_alerts)})


# ── CVE Stats ────────────────────────────────────────────────────────

@android_forensics_ai_bp.route('/cve-stats', methods=['GET'])
@login_required
def cve_stats():
    return jsonify({'ok': True, 'stats': _cve().get_stats()})