"""3D Globe view and flow extraction."""
import json
import logging
import traceback

from django.shortcuts import render

from analysis.constants import GLOBE_MAX_ROWS
from analysis.models import AnalysisRun
from .helpers import _extract_lat, _extract_lon

logger = logging.getLogger(__name__)


def _extract_flows_from_df(df, MAX_ROWS):
    """Extract flow data (lat/lon/TLS/bytes) from a Polars DataFrame."""
    from analysis.geoip import lookup as geo_lookup
    from analysis.type_classifier import TLS_HEX_MAP

    src_lat_col = next((c for c in df.columns if c.lower() in (
        ':ips.latd', ':ips_latd', ':ips_lat', 'src_latitude', 'src_lat', 'source_latitude', 'source_lat')), None)
    src_lon_col = next((c for c in df.columns if c.lower() in (
        ':ips.lond', ':ips_lond', ':ips_lon', 'src_longitude', 'src_lon', 'source_longitude', 'source_lon')), None)
    dst_lat_col = next((c for c in df.columns if c.lower() in (
        ':ipd.latd', ':ipd_latd', ':ipd_lat', 'dst_latitude', 'dst_lat', 'dest_latitude', 'dest_lat')), None)
    dst_lon_col = next((c for c in df.columns if c.lower() in (
        ':ipd.lond', ':ipd_lond', ':ipd_lon', 'dst_longitude', 'dst_lon', 'dest_longitude', 'dest_lon')), None)
    src_ip_col = next((c for c in df.columns if 'src_ip' in c.lower() or ':ips' in c.lower()), None)
    dst_ip_col = next((c for c in df.columns if 'dst_ip' in c.lower() or ':ipd' in c.lower()), None)
    tls_col = next((c for c in df.columns if c.lower() in ('0ver', 'tls_version', '_tlsver', 'version')), None)
    bytes_col = next((c for c in df.columns if c.lower() in ('8ack', '8pak', '8byt', 'bytes_sent', 'bytes')), None)
    time_col = next((c for c in df.columns if c.lower() in (
        'timestamp', 'time', '_time', 'ts', 'datetime', 'epoch',
        'created_at', 'created', 'event_time')), None)

    flows = []
    for row in df.iter_rows(named=True):
        slat = slon = dlat = dlon = None

        if src_lat_col and src_lon_col:
            slat = _extract_lat(row.get(src_lat_col))
            slon = _extract_lon(row.get(src_lon_col))
        if dst_lat_col and dst_lon_col:
            dlat = _extract_lat(row.get(dst_lat_col))
            dlon = _extract_lon(row.get(dst_lon_col))

        if (slat is None or slon is None) and src_ip_col:
            sg = geo_lookup(str(row.get(src_ip_col, '')))
            if sg:
                slat, slon = sg['lat'], sg['lon']
        if (dlat is None or dlon is None) and dst_ip_col:
            dg = geo_lookup(str(row.get(dst_ip_col, '')))
            if dg:
                dlat, dlon = dg['lat'], dg['lon']

        if slat is None or slon is None or dlat is None or dlon is None:
            continue

        tls_val = str(row.get(tls_col or '', ''))
        # Check hex format first (e.g. "03 03" or "0303")
        hex_key = tls_val.replace(' ', '').strip()
        if hex_key in TLS_HEX_MAP:
            tls_label = TLS_HEX_MAP[hex_key]
        elif '1.3' in tls_val:
            tls_label = 'TLSv1.3'
        elif '1.2' in tls_val:
            tls_label = 'TLSv1.2'
        else:
            tls_label = 'other'

        bs = float(row.get(bytes_col or '', 0) or 0)

        # Extract timestamp for animation ordering
        t_raw = row.get(time_col) if time_col else None
        t_val = None
        if t_raw is not None:
            if isinstance(t_raw, (int, float)):
                t_val = float(t_raw)
            elif hasattr(t_raw, 'timestamp'):
                t_val = t_raw.timestamp()
            else:
                try:
                    t_val = float(t_raw)
                except (ValueError, TypeError):
                    pass

        flows.append({
            'slat': slat, 'slon': slon,
            'dlat': dlat, 'dlon': dlon,
            'tls': tls_label,
            'bytes': bs,
            'time': t_val,
        })

    # Normalize timestamps to 0-1 and sort flows by time
    timed = [f for f in flows if f['time'] is not None]
    if len(timed) > 1:
        timed.sort(key=lambda f: f['time'])
        t_min, t_max = timed[0]['time'], timed[-1]['time']
        t_range = max(t_max - t_min, 1)
        for f in timed:
            f['time'] = (f['time'] - t_min) / t_range
        # Untimed flows get evenly spaced after timed ones
        untimed = [f for f in flows if f['time'] is None]
        for i, f in enumerate(untimed):
            f['time'] = min(1.0, (len(timed) + i) / max(len(flows) - 1, 1))
        flows = timed + untimed
    else:
        for i, f in enumerate(flows):
            f['time'] = i / max(len(flows) - 1, 1)

    return flows


def globe_view(request):
    from analysis.geoip import lookup as geo_lookup
    from analysis.data_loader import load_csv_directory, load_from_db
    from analysis.models import AnalysisRun
    from analysis.session_store import SessionStore

    MAX_ROWS_PER_RUN = GLOBE_MAX_ROWS  # Limit per run to prevent browser overload

    # Check for ?data=dataset_id parameter (direct data loading from session store)
    data_param = request.GET.get('data')
    embed_mode = request.GET.get('embed') == '1'
    template_name = 'tianxuan/globe_embed.html' if embed_mode else 'tianxuan/globe.html'

    if data_param:
        store = SessionStore()
        flows = []
        all_runs = list(AnalysisRun.objects.filter(status='completed').order_by('-id').values('display_id', 'csv_glob'))
        try:
            entry = store.get_dataset(data_param)
            if entry is not None:
                lf = entry['lazyframe']
                df = lf.head(MAX_ROWS_PER_RUN).collect(streaming=True)
                flows = _extract_flows_from_df(df, MAX_ROWS_PER_RUN)
        except Exception as exc:
            logging.getLogger(__name__).warning(f'globe ?data= param failed: {exc}')
        return render(request, template_name, {
            'geo_flows': json.dumps(flows),
            'all_runs': all_runs,
            'selected_ids': [],
        })

    # Get all completed runs for the selector
    all_runs = list(AnalysisRun.objects.filter(status='completed').order_by('-id').values('display_id', 'csv_glob'))

    # Get selected run IDs from query params (multiple checkboxes)
    selected_ids_list = request.GET.getlist('runs')
    selected_runs = []
    if selected_ids_list:
        for sid in selected_ids_list:
            sid = sid.strip()
            if sid.isdigit():
                try:
                    r = AnalysisRun.objects.get(display_id=int(sid), status='completed')
                    selected_runs.append(r)
                except AnalysisRun.DoesNotExist:
                    pass

    # Check if user-selected runs have valid files; if not, fall through
    all_selected_failed = bool(selected_ids_list)
    fallback_run = None

    if not selected_runs:
        all_selected_failed = False

    flows = []

    for run in selected_runs:
        try:
            lf = None
            if run.sqlite_table:
                lf = load_from_db(run.sqlite_table)
            if lf is None:
                lf, schema, rc, fc, mem = load_csv_directory(run.csv_glob)
            df = lf.head(MAX_ROWS_PER_RUN).collect(streaming=True)
            flows.extend(_extract_flows_from_df(df, MAX_ROWS_PER_RUN))
            all_selected_failed = False
            break  # stop after first successful load
        except Exception as exc:
            logger.warning('[GLOBE] Failed to load run %s (%s): %s', run.display_id, run.csv_glob, exc)
            continue

    if all_selected_failed or not flows:
        # Fallback: latest completed run with valid data (SQLite or CSV)
        for latest in AnalysisRun.objects.filter(status='completed').order_by('-id'):
            if not latest.csv_glob and not latest.sqlite_table:
                continue
            try:
                lf = None
                if latest.sqlite_table:
                    lf = load_from_db(latest.sqlite_table)
                if lf is None:
                    if not latest.csv_glob:
                        continue
                    lf, schema, rc, fc, mem = load_csv_directory(latest.csv_glob)
                df = lf.head(MAX_ROWS_PER_RUN).collect(streaming=True)
                flows.extend(_extract_flows_from_df(df, MAX_ROWS_PER_RUN))
                selected_runs = [latest]
                break
            except Exception:
                logger.error("unknown failed: {}".format(traceback.format_exc()))
                continue

    resp = render(request, 'tianxuan/globe.html', {
        'geo_flows': json.dumps(flows),
        'all_runs': all_runs,
        'selected_ids': [r.display_id for r in selected_runs],
    })
    resp['Cache-Control'] = 'no-cache, no-store, must-revalidate'
    resp['Pragma'] = 'no-cache'
    resp['Expires'] = '0'
    return resp
