import pandas as pd
import numpy as np
import requests
import uvicorn
import asyncio
import re
import uuid
import json
import io
import zipfile
from datetime import timedelta
from fastapi import FastAPI, Body, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse, Response
from huggingface_hub import HfApi, hf_hub_download
from typing import Optional, List

# ==========================================
# 1. CONFIGURATION & SECRETS
# ==========================================
SB_URL = "https://txmakjdjvtwvdbwkeopd.supabase.co"
SB_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InR4bWFramRqdnR3dmRid2tlb3BkIiwicm9sZSI6InNlcnZpY2Vfcm9sZSIsImlhdCI6MTc4MTQ4ODU2OSwiZXhwIjoyMDk3MDY0NTY5fQ.OQTQ9M6SlCYGRZdQgmFQVHfwUhVKUQ0iBcLuitrHdok"
HEADERS = {
    "apikey": SB_KEY,
    "Authorization": f"Bearer {SB_KEY}",
    "Content-Type": "application/json",
    "Prefer": "return=representation"
}

HF_TOKEN_PART_1 = "hf_EzaiwOHiejWFLN"
HF_TOKEN_PART_2 = "otvQqTPUTksQQfYxrbJg"
HF_TOKEN = HF_TOKEN_PART_1 + HF_TOKEN_PART_2
HF_REPO_ID = "topsecrettraders/DepthChain"

hf_api = HfApi(token=HF_TOKEN)
app = FastAPI()

SESSION_CACHE = {}

# ==========================================
# 2. SUPABASE HELPER FUNCTIONS
# ==========================================
def sb_rpc(func_name: str, params: dict = None):
    if params is None: params = {}
    try:
        url = f"{SB_URL}/rest/v1/rpc/{func_name}"
        r = requests.post(url, headers=HEADERS, json=params)
        r.raise_for_status()
        return r.json()
    except Exception as e:
        err_msg = r.text if 'r' in locals() and hasattr(r, 'text') else str(e)
        print(f"DB Error RPC ({func_name}): {err_msg}")
        return []

# ==========================================
# 3. API ROUTES - METADATA & UI
# ==========================================
@app.get("/")
def home():
    return HTMLResponse(HTML_TEMPLATE)

@app.get("/api/dates")
def get_dates():
    data = sb_rpc("get_available_dates", {"p_limit": 50})
    if isinstance(data, list):
        return JSONResponse([row['available_date'] for row in data if 'available_date' in row])
    return JSONResponse([])

@app.get("/api/roots")
def get_roots(date: str):
    data = sb_rpc("get_roots_for_date", {"p_date": date, "p_q": ""})
    if isinstance(data, list):
        return JSONResponse([row['root_name'] for row in data if 'root_name' in row])
    return JSONResponse([])

@app.get("/api/expiries")
def get_expiries(date: str, root: str):
    data = sb_rpc("get_expiries_for_root", {"p_date": date, "p_root": root.upper()})
    if isinstance(data, list):
        return JSONResponse([row['expiry_date'] for row in data if 'expiry_date' in row])
    return JSONResponse([])

# ==========================================
# 4. HUGGING FACE EXPLORER API
# ==========================================
@app.get("/api/hf_files")
def list_hf_files():
    try:
        files = hf_api.list_repo_files(repo_id=HF_REPO_ID, repo_type="dataset")
        sessions = {}
        
        for f in files:
            if f.startswith("Sessions/") and f.endswith("data.parquet"):
                parts = f.split("/")
                if len(parts) >= 4:
                    session_name = parts[1]
                    dataset_name = parts[2]
                    
                    if session_name not in sessions:
                        sessions[session_name] = []
                    if dataset_name not in sessions[session_name]:
                        sessions[session_name].append(dataset_name)
                        
        # Sort sessions and datasets
        sorted_sessions = {k: sorted(v, reverse=True) for k, v in sorted(sessions.items(), reverse=True)}
        return JSONResponse({"status": "success", "sessions": sorted_sessions})
    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)})

@app.get("/api/hf_metadata")
def get_hf_metadata(session_name: str, dataset_name: str):
    try:
        meta_path = f"Sessions/{session_name}/{dataset_name}/metadata.json"
        local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=meta_path, repo_type="dataset", token=HF_TOKEN)
        with open(local_path, "r") as f:
            metadata = json.load(f)
        return JSONResponse({"status": "success", "metadata": metadata})
    except Exception as e:
        return JSONResponse({"status": "error", "message": "Metadata not found or error fetching."})

@app.get("/api/view_parquet")
def view_parquet(session_name: str, dataset_name: str):
    try:
        file_path = f"Sessions/{session_name}/{dataset_name}/data.parquet"
        local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=file_path, repo_type="dataset", token=HF_TOKEN)
        df = pd.read_parquet(local_path)
        # Limit rows to 1000 to prevent browser crash, it's just a viewer
        df_preview = df.head(1000)
        return JSONResponse({"status": "success", "columns": df_preview.columns.tolist(), "data": df_preview.to_dict(orient="records")})
    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)})

@app.get("/api/download_session_zip")
def download_session_zip(session_name: str):
    try:
        folder_path = f"Sessions/{session_name}"
        files = hf_api.list_repo_files(repo_id=HF_REPO_ID, repo_type="dataset")
        folder_files = [f for f in files if f.startswith(folder_path + "/")]

        if not folder_files:
            return JSONResponse({"status": "error", "message": "No files found for this session."})

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
            for f_path in folder_files:
                local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=f_path, repo_type="dataset", token=HF_TOKEN)
                # Keep the dataset folder structure inside the zip
                arcname = f_path.replace(f"{folder_path}/", "")
                zip_file.write(local_path, arcname=arcname)

        return Response(
            content=zip_buffer.getvalue(),
            media_type="application/zip",
            headers={"Content-Disposition": f'attachment; filename="{session_name}_Full.zip"'}
        )
    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)})

@app.get("/api/download_session_parquets_zip")
def download_session_parquets_zip(session_name: str):
    try:
        folder_path = f"Sessions/{session_name}"
        files = hf_api.list_repo_files(repo_id=HF_REPO_ID, repo_type="dataset")
        
        # Only grab the data.parquet files for this session
        parquet_files = [f for f in files if f.startswith(folder_path + "/") and f.endswith("data.parquet")]

        if not parquet_files:
            return JSONResponse({"status": "error", "message": "No Parquet files found for this session."})

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
            for f_path in parquet_files:
                local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=f_path, repo_type="dataset", token=HF_TOKEN)
                
                # Extract dataset name (e.g. 2026-07-10_NIFTY_2026-07-14)
                # Structure: Sessions / {session_name} / {dataset_name} / data.parquet
                parts = f_path.split("/")
                if len(parts) >= 4:
                    dataset_name = parts[2]
                    # Rename the file inside the zip to match requested format
                    arcname = f"{dataset_name}_data.parquet"
                    zip_file.write(local_path, arcname=arcname)

        return Response(
            content=zip_buffer.getvalue(),
            media_type="application/zip",
            headers={"Content-Disposition": f'attachment; filename="{session_name}_Parquets_Only.zip"'}
        )
    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)})

# ==========================================
# 5. STEP 1: FETCH & INSPECT
# ==========================================
@app.post("/api/inspect_data")
async def inspect_data(
    date: str = Body(...),
    root: str = Body(...),
    expiry: str = Body(...),
    start_time: str = Body("09:15"),
    end_time: str = Body("15:30"),
    atm_range: int = Body(10),
    required_vars: List[str] = Body(["P", "CB", "PS", "CS", "PB"])
):
    root = root.upper()
    t_start_dt = pd.to_datetime(f"{date} {start_time}:00")
    t_end_dt = pd.to_datetime(f"{date} {end_time}:00")

    inst_data = sb_rpc("get_instruments_for_expiry", {"p_date": date, "p_root": root, "p_expiry": expiry})
    instruments = [r['instrument_name'] for r in inst_data if 'instrument_name' in r]

    if not instruments:
        return JSONResponse({"error": f"No data found in Supabase for {root} {expiry} on {date}"})

    spot_keys =[k for k in instruments if k.endswith('-INDEX') or k.endswith('-EQ')]
    fut_keys =[k for k in instruments if 'FUT' in k]
    ref_instrument = spot_keys[0] if spot_keys else (fut_keys[0] if fut_keys else instruments[0])

    strike_map = {}
    for inst in instruments:
        if inst.endswith('CE') or inst.endswith('PE'):
            clean_inst = inst.split(":")[-1]
            if clean_inst.startswith(root):
                match = re.search(r'(\d+(?:\.\d+)?)(?:CE|PE)$', clean_inst[len(root)+5:])
                if match:
                    strike_map[inst] = float(match.group(1))
                    continue
            matches = re.findall(r'(\d+(?:\.\d+)?)(?=CE$|PE$)', clean_inst)
            if matches: strike_map[inst] = float(matches[-1])

    base_params = {
        "p_root": root, "p_expiry": expiry, "p_ref_instrument": ref_instrument,
        "p_atm_range": atm_range, "p_strike_map": strike_map
    }

    async def fetch_chunk(start_dt, end_dt, retries=3):
        params = base_params.copy()
        params["p_start_time"] = start_dt.strftime("%Y-%m-%d %H:%M:%S")
        params["p_end_time"] = end_dt.strftime("%Y-%m-%d %H:%M:%S")
        loop = asyncio.get_event_loop()
        for attempt in range(retries):
            try:
                return await loop.run_in_executor(None, sb_rpc, "get_advanced_chart_data", params)
            except Exception as e:
                if attempt == retries - 1: return[]
                await asyncio.sleep(1)

    tasks =[]
    curr = t_start_dt
    while curr <= t_end_dt:
        nxt = curr + timedelta(hours=1) - timedelta(seconds=1)
        if nxt > t_end_dt: nxt = t_end_dt
        tasks.append(fetch_chunk(curr, nxt))
        curr += timedelta(hours=1)

    chunk_results = await asyncio.gather(*tasks)
    
    all_records =[]
    for res in chunk_results:
        if isinstance(res, list): all_records.extend(res)

    if not all_records:
        return JSONResponse({"error": "No valid records fetched. Data might be empty for this date/time."})

    df = pd.DataFrame(all_records)
    df.rename(columns={"tick_ts": "timestamp", "price": "P", "cb": "CB", "cs": "CS", "pb": "PB", "ps": "PS", "b_qty": "B", "s_qty": "S", "vol": "V", "oi": "OI"}, inplace=True)
    df['timestamp'] = pd.to_datetime(df['timestamp'])
    df.drop_duplicates(subset=['timestamp'], inplace=True)
    df.set_index('timestamp', inplace=True)
    df.sort_index(inplace=True)

    df = df.loc[(df.index >= t_start_dt) & (df.index <= t_end_dt)]
    if df.empty:
        return JSONResponse({"error": "Data exists, but completely outside your specified Start/End Time."})

    available_vars =[v for v in required_vars if v in df.columns]
    df = df[available_vars].copy()

    actual_min = df.index.min()
    actual_max = df.index.max()
    full_idx = pd.date_range(start=actual_min, end=actual_max, freq='10s')
    
    raw_ticks = len(df)
    expected_ticks = len(full_idx)
    missing_ts = full_idx.difference(df.index).to_pydatetime()
    
    missing_blocks =[]
    if len(missing_ts) > 0:
        start_block = missing_ts[0]
        prev = missing_ts[0]
        for ts in missing_ts[1:]:
            if (ts - prev).total_seconds() > 10.5:
                missing_blocks.append({
                    "start": start_block.strftime('%Y-%m-%d %H:%M:%S'),
                    "end": prev.strftime('%Y-%m-%d %H:%M:%S'),
                    "count": int((prev - start_block).total_seconds()/10) + 1
                })
                start_block = ts
            prev = ts
        missing_blocks.append({
            "start": start_block.strftime('%Y-%m-%d %H:%M:%S'),
            "end": prev.strftime('%Y-%m-%d %H:%M:%S'),
            "count": int((prev - start_block).total_seconds()/10) + 1
        })

    df_reindexed = df.reindex(full_idx)
    null_details_existing = {}
    for col in available_vars:
        na_times = df_reindexed[df_reindexed[col].isna() | (df_reindexed[col] == 0)].index.strftime('%H:%M:%S').tolist()
        null_details_existing[col] = {"count": len(na_times), "times": na_times[:50]} 

    session_id = str(uuid.uuid4())
    SESSION_CACHE[session_id] = df.copy()

    return JSONResponse({
        "status": "success",
        "session_id": session_id,
        "stats": {
            "expected_ticks": expected_ticks,
            "found_ticks": raw_ticks,
            "missing_time_ticks": len(missing_ts),
            "missing_blocks": missing_blocks,
            "null_breakdown": null_details_existing,
            "variables": available_vars,
            "original_start": actual_min.strftime('%Y-%m-%d %H:%M:%S'),
            "original_end": actual_max.strftime('%Y-%m-%d %H:%M:%S')
        },
        "meta": {
            "date": date, "root": root, "expiry": expiry,
            "start": start_time, "end": end_time, "atm_range": atm_range
        }
    })

# ==========================================
# 6. STEP 2: CLEAN & ARCHIVE
# ==========================================
@app.post("/api/clean_and_upload")
async def clean_and_upload(
    session_id: str = Body(...),
    session_name: str = Body(...),
    fill_method: str = Body("ffill"),
    trim_start: str = Body(None),
    trim_end: str = Body(None),
    metadata: dict = Body(...)
):
    if session_id not in SESSION_CACHE:
        return JSONResponse({"error": "Session expired or invalid. Please Fetch & Inspect again."})

    # Validate session name
    clean_session_name = re.sub(r'[^a-zA-Z0-9_\-]', '_', session_name)
    if not clean_session_name:
        clean_session_name = "Default_Session"

    df = SESSION_CACHE[session_id].copy()

    t_start = pd.to_datetime(trim_start) if trim_start else df.index.min()
    t_end = pd.to_datetime(trim_end) if trim_end else df.index.max()
    df = df.loc[(df.index >= t_start) & (df.index <= t_end)]

    full_idx = pd.date_range(start=t_start, end=t_end, freq='10s')
    df = df.reindex(full_idx)
    df.replace(0, np.nan, inplace=True)

    if fill_method == "ffill":
        df_clean = df.ffill().fillna(0)  
    elif fill_method == "interpolate":
        df_clean = df.interpolate(method='linear').fillna(0)
    elif fill_method == "zero":
        df_clean = df.fillna(0)
    else:
        df_clean = df.ffill().fillna(0)

    df_clean.reset_index(inplace=True)
    df_clean.rename(columns={"index": "timestamp"}, inplace=True)
    df_clean['timestamp'] = df_clean['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S')

    metadata["session_name"] = clean_session_name
    metadata["imputation_applied"] = fill_method
    metadata["trim_start_applied"] = trim_start
    metadata["trim_end_applied"] = trim_end
    metadata["total_final_rows"] = len(df_clean)

    date = metadata["date"]
    root = metadata["root"]
    expiry = metadata["expiry"]
    
    target_ds_name = f"{date}_{root}_{expiry}"
    
    # Handle duplicates inside the specific session folder
    try:
        files = hf_api.list_repo_files(repo_id=HF_REPO_ID, repo_type="dataset")
        session_files = [f for f in files if f.startswith(f"Sessions/{clean_session_name}/")]
        
        final_ds_name = target_ds_name
        version = 1
        while any(f.startswith(f"Sessions/{clean_session_name}/{final_ds_name}/") for f in session_files):
            version += 1
            final_ds_name = f"{target_ds_name}_v{version}"
            
        hf_folder = f"Sessions/{clean_session_name}/{final_ds_name}"
    except Exception as e:
        hf_folder = f"Sessions/{clean_session_name}/{target_ds_name}"

    try:
        parquet_buffer = io.BytesIO()
        df_clean.to_parquet(parquet_buffer, index=False)
        parquet_buffer.seek(0)
        
        json_buffer = io.BytesIO(json.dumps(metadata, indent=4).encode('utf-8'))
        json_buffer.seek(0)

        hf_api.upload_file(path_or_fileobj=parquet_buffer, path_in_repo=f"{hf_folder}/data.parquet", repo_id=HF_REPO_ID, repo_type="dataset")
        hf_api.upload_file(path_or_fileobj=json_buffer, path_in_repo=f"{hf_folder}/metadata.json", repo_id=HF_REPO_ID, repo_type="dataset")
        
        del SESSION_CACHE[session_id]
        return JSONResponse({"status": "success", "message": f"Dataset safely archived to {hf_folder}!"})
    except Exception as e:
        return JSONResponse({"error": f"Upload Failed: {str(e)}"})


# ==========================================
# 7. FRONTEND TEMPLATE (LIGHT MODE)
# ==========================================
HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>DepthChain Data Engineering</title>
    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap" rel="stylesheet">
    <style>
        :root { 
            --bg: #f8fafc; 
            --panel: #ffffff; 
            --border: #e2e8f0; 
            --text-main: #0f172a; 
            --text-muted: #64748b; 
            --accent: #2563eb; 
            --accent-hover: #1d4ed8;
            --success: #16a34a; 
            --success-bg: #dcfce7;
            --danger: #dc2626; 
            --danger-bg: #fee2e2;
            --warn: #d97706;
            --warn-bg: #fef3c7;
        }
        body { background: var(--bg); color: var(--text-main); font-family: 'Inter', sans-serif; margin: 0; padding: 30px; font-size: 13px; }
        .header-title { font-size: 24px; font-weight: 700; color: var(--text-main); margin-bottom: 20px; display: flex; align-items: center; gap: 10px; }
        .grid { display: grid; grid-template-columns: 1fr 1.3fr; gap: 24px; max-width: 1400px; margin: 0 auto; }
        .card { background: var(--panel); border: 1px solid var(--border); border-radius: 12px; padding: 24px; display: flex; flex-direction: column; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05); }
        h2 { margin-top: 0; font-size: 16px; border-bottom: 1px solid var(--border); padding-bottom: 12px; color: var(--text-main); font-weight: 600; }
        
        label { display: block; margin-bottom: 6px; color: var(--text-muted); font-size: 11px; font-weight: 600; text-transform: uppercase; letter-spacing: 0.5px; }
        select, input { width: 100%; padding: 10px 14px; margin-bottom: 16px; background: #f8fafc; border: 1px solid var(--border); color: var(--text-main); border-radius: 8px; font-size: 13px; font-weight: 500; box-sizing: border-box; transition: all 0.2s; }
        select:focus, input:focus { outline: none; border-color: var(--accent); background: #fff; box-shadow: 0 0 0 3px rgba(37,99,235,0.1); }
        
        button { background: var(--accent); color: white; border: none; padding: 12px; border-radius: 8px; font-weight: 600; font-size: 13px; cursor: pointer; transition: 0.2s; width: 100%; display: flex; align-items: center; justify-content: center; gap: 8px; }
        button:hover { background: var(--accent-hover); transform: translateY(-1px); }
        button:disabled { background: var(--border); color: var(--text-muted); cursor: not-allowed; transform: none; }
        
        .btn-success { background: var(--success); }
        .btn-success:hover { background: #15803d; }
        .btn-outline { background: transparent; border: 1px solid var(--border); color: var(--text-main); width: auto; padding: 8px 16px; font-size: 12px; }
        .btn-outline:hover { background: #f1f5f9; border-color: #cbd5e1; color: var(--text-main); }
        .btn-warn { background: transparent; border: 1px solid var(--warn); color: var(--warn); padding: 4px 10px; font-size: 11px; border-radius: 6px; width: auto; }
        .btn-warn:hover { background: var(--warn-bg); }
        
        .row { display: flex; gap: 16px; }
        .col { flex: 1; }
        
        .stat-box { background: #f8fafc; border: 1px solid var(--border); border-radius: 8px; padding: 16px; margin-bottom: 16px; }
        .stat-item { display: flex; justify-content: space-between; margin-bottom: 10px; border-bottom: 1px dashed #cbd5e1; padding-bottom: 6px; align-items: center; }
        .stat-item:last-child { border-bottom: none; margin-bottom: 0; padding-bottom: 0; }
        .stat-val { font-family: 'JetBrains Mono', monospace; font-size: 13px; font-weight: 700; }
        .val-good { color: var(--success); }
        .val-bad { color: var(--danger); }
        .val-warn { color: var(--warn); }
        
        /* Session Explorer Styling */
        .hf-list { flex: 1; overflow-y: auto; max-height: 600px; padding-right: 5px; }
        .session-group { border: 1px solid var(--border); border-radius: 8px; margin-bottom: 16px; overflow: hidden; background: #fff; }
        .session-header { background: #f1f5f9; padding: 12px 16px; font-weight: 600; display: flex; justify-content: space-between; align-items: center; border-bottom: 1px solid var(--border); color: var(--text-main); flex-wrap: wrap; gap: 10px; }
        .session-datasets { padding: 8px 16px; }
        .dataset-item { padding: 10px 0; border-bottom: 1px solid var(--border); display: flex; justify-content: space-between; align-items: center; }
        .dataset-item:last-child { border-bottom: none; }
        .dataset-name { font-weight: 500; color: var(--text-muted); font-size: 12.5px; }
        
        .hf-btns { display: flex; gap: 8px; flex-wrap: wrap; justify-content: flex-end; }
        .hf-btn { cursor: pointer; font-size: 11px; padding: 6px 10px; background: #fff; border: 1px solid var(--border); color: var(--text-main); border-radius: 6px; font-weight: 500; }
        .hf-btn:hover { background: #f8fafc; border-color: var(--accent); color: var(--accent); }
        
        .btn-download-session { background: var(--success); color: white; padding: 6px 12px; font-size: 11px; border: none; border-radius: 6px; width: auto; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
        .btn-download-session:hover { filter: brightness(0.9); }

        .toast { position: fixed; bottom: 30px; right: 30px; background: var(--success); color: white; padding: 16px 24px; border-radius: 8px; display: none; box-shadow: 0 10px 15px -3px rgba(0,0,0,0.1); font-weight: 600; z-index: 9999; }
        
        details { margin-bottom: 5px; font-size: 12px; }
        summary { cursor: pointer; color: var(--accent); font-weight: 600; margin-bottom: 5px; outline: none; }
        .inspect-times { font-family: 'JetBrains Mono', monospace; font-size: 11px; color: var(--danger); background: var(--danger-bg); padding: 8px; border-radius: 6px; border: 1px solid #fca5a5; max-height: 100px; overflow-y: auto; margin-top: 4px; }

        .modal-bg { position: fixed; top: 0; left: 0; width: 100%; height: 100%; background: rgba(15, 23, 42, 0.6); backdrop-filter: blur(4px); display: none; justify-content: center; align-items: center; z-index: 1000; }
        .modal-content { background: var(--panel); border-radius: 12px; width: 95%; max-width: 1100px; height: 85vh; display: flex; flex-direction: column; box-shadow: 0 25px 50px -12px rgba(0, 0, 0, 0.25); }
        .modal-header { padding: 20px; border-bottom: 1px solid var(--border); display: flex; justify-content: space-between; align-items: center; font-weight: 700; font-size: 16px; color: var(--text-main); }
        .modal-body { flex: 1; padding: 0; overflow: auto; background: #f8fafc; }
        .close-btn { background: #f1f5f9; border: none; color: var(--text-muted); width: 32px; height: 32px; border-radius: 50%; display: flex; align-items: center; justify-content: center; font-size: 16px; cursor: pointer; transition: 0.2s; padding: 0; }
        .close-btn:hover { background: #e2e8f0; color: var(--danger); }
        
        table.data-table { width: 100%; border-collapse: collapse; font-family: 'JetBrains Mono', monospace; font-size: 12px; text-align: right; background: #fff;}
        table.data-table th { background: #f1f5f9; color: var(--text-muted); padding: 12px 16px; border-bottom: 1px solid var(--border); position: sticky; top: 0; font-weight: 600; text-transform: uppercase; font-size: 11px; z-index: 10; }
        table.data-table td { padding: 10px 16px; border-bottom: 1px solid var(--border); color: var(--text-main); }
        table.data-table tr:hover { background: #f8fafc; }
    </style>
</head>
<body>

<div class="header-title">
    ⚡ DepthChain Data Studio
</div>

<div class="grid">
    <!-- LEFT COLUMN -->
    <div style="display: flex; flex-direction: column; gap: 24px;">
        
        <!-- SESSION DEFINITION -->
        <div class="card">
            <h2>1. Define Active Session</h2>
            <label>Session Name (Creates a master folder)</label>
            <input type="text" id="inpSessionName" placeholder="e.g. Backtest_July_Run" style="margin-bottom:0;">
            <p style="font-size: 11px; color: var(--text-muted); margin-top: 6px; margin-bottom:0;">All datasets extracted will be saved inside this session folder.</p>
        </div>

        <div class="card">
            <h2>2. Extract & Inspect Data</h2>
            <div class="row">
                <div class="col"><label>Date</label><select id="selDate" onchange="App.loadRoots()"><option>Loading...</option></select></div>
                <div class="col"><label>Symbol</label><select id="selRoot" onchange="App.loadExpiries()"><option>Loading...</option></select></div>
            </div>
            <label>Expiry Date</label><select id="selExpiry"><option>Loading...</option></select>
            
            <div class="row">
                <div class="col"><label>ATM Range (± N)</label><input type="number" id="inpAtm" value="10"></div>
                <div class="col"><label>Variables</label><input type="text" id="inpVars" value="P,CB,PS,CS,PB"></div>
            </div>
            <div class="row">
                <div class="col"><label>Strict Start</label><input type="time" id="inpStart" value="09:15"></div>
                <div class="col"><label>Strict End</label><input type="time" id="inpEnd" value="15:30"></div>
            </div>
            
            <button id="btnInspect" onclick="App.inspectData()">🔍 Fetch & Analyze Quality</button>
        </div>

        <div class="card" id="cardArchive" style="opacity: 0.5; pointer-events: none;">
            <h2>3. Clean & Archive to Session</h2>
            <div class="stat-box" id="statsBox">
                <div style="text-align:center; color:var(--text-muted); font-style:italic; padding: 10px 0;">Run Inspection first...</div>
            </div>
            
            <label>Missing Data Handling Strategy</label>
            <select id="selFill">
                <option value="ffill">Forward Fill (Fixes Nulls & Zeroes with Previous Values) - REC.</option>
                <option value="interpolate">Linear Interpolation (Average between gaps)</option>
                <option value="zero">Fill with Zeroes</option>
            </select>
            
            <button id="btnUpload" class="btn-success" onclick="App.uploadData()">☁️ Clean & Archive to Session</button>
        </div>
    </div>

    <!-- RIGHT COLUMN -->
    <div class="card">
        <div style="display:flex; justify-content:space-between; align-items:center; border-bottom: 1px solid var(--border); padding-bottom: 12px; margin-bottom: 16px;">
            <h2 style="border:none; padding:0; margin:0;">Hugging Face Repository Explorer</h2>
            <button class="btn-outline" onclick="App.loadHFFolders()">🔄 Refresh</button>
        </div>
        
        <div class="hf-list" id="hfList">
            <div style="text-align:center; padding:40px; color:var(--text-muted);">Loading sessions...</div>
        </div>
        
        <div class="stat-box" id="metaBox" style="margin-top: 20px; display: none; background: #1e293b; color: #f8fafc; border:none;">
            <div style="display:flex; justify-content:space-between; align-items:center; margin-bottom:12px;">
                <h3 style="margin:0; font-size:13px; color:#cbd5e1;">Metadata Viewer</h3>
                <button style="width:auto; padding:4px 8px; font-size:11px; background:#334155; color:white; border:none; border-radius:4px;" onclick="document.getElementById('metaBox').style.display='none'">Close</button>
            </div>
            <pre id="metaContent" style="font-family:'JetBrains Mono', monospace; font-size:11.5px; color:#a5b4fc; margin:0; white-space:pre-wrap;"></pre>
        </div>
    </div>
</div>

<div class="toast" id="toast">✅ Success!</div>

<!-- DATA VIEWER MODAL -->
<div class="modal-bg" id="dataModal">
    <div class="modal-content">
        <div class="modal-header">
            <span id="modalTitle">Data Viewer</span>
            <button class="close-btn" onclick="document.getElementById('dataModal').style.display='none'">✖</button>
        </div>
        <div class="modal-body" id="modalBody">
            <div style="text-align:center; padding: 50px; color:var(--text-muted);">Loading data table...</div>
        </div>
    </div>
</div>

<script>
const App = {
    currentSession: null,
    currentMeta: null,
    trimStart: null,
    trimEnd: null,
    originalStartStr: null,
    originalEndStr: null,
    missingTimeTicks: 0,

    async init() {
        // Set default session name to today's date
        const today = new Date().toISOString().split('T')[0];
        document.getElementById('inpSessionName').value = `Session_${today}`;

        await this.loadDates();
        await this.loadHFFolders();
    },

    showToast(msg) {
        const t = document.getElementById('toast');
        t.innerText = msg;
        t.style.display = 'block';
        setTimeout(() => t.style.display = 'none', 4000);
    },

    async loadDates() {
        const sel = document.getElementById('selDate');
        sel.innerHTML = '<option>Fetching...</option>';
        try {
            const res = await fetch('/api/dates');
            const dates = await res.json();
            if (dates && dates.length > 0) {
                sel.innerHTML = dates.map(d => `<option value="${d}">${d}</option>`).join('');
                this.loadRoots();
            } else {
                sel.innerHTML = '<option value="">No dates available</option>';
            }
        } catch(e) {
            sel.innerHTML = '<option value="">Database Connection Error</option>';
        }
    },

    async loadRoots() {
        const date = document.getElementById('selDate').value;
        const sel = document.getElementById('selRoot');
        if(!date) return;
        
        // Save current symbol selection if it exists so we can preserve it across date changes
        const currentSymbol = sel.value;
        
        sel.innerHTML = '<option>Fetching...</option>';
        try {
            const res = await fetch(`/api/roots?date=${date}`);
            const roots = await res.json();
            if (roots && roots.length > 0) {
                sel.innerHTML = roots.map(r => `<option value="${r}">${r}</option>`).join('');
                
                // If the previously selected symbol exists for this new date, auto-select it!
                if (currentSymbol && roots.includes(currentSymbol)) {
                    sel.value = currentSymbol;
                }
                
                this.loadExpiries();
            } else {
                sel.innerHTML = '<option value="">No symbols found</option>';
                document.getElementById('selExpiry').innerHTML = '<option value="">-</option>';
            }
        } catch(e) {
            sel.innerHTML = '<option value="">Error</option>';
        }
    },

    async loadExpiries() {
        const date = document.getElementById('selDate').value;
        const root = document.getElementById('selRoot').value;
        const sel = document.getElementById('selExpiry');
        if(!root || root === 'Fetching...') return;
        
        sel.innerHTML = '<option>Fetching...</option>';
        try {
            const res = await fetch(`/api/expiries?date=${date}&root=${root}`);
            const exps = await res.json();
            if (exps && exps.length > 0) {
                sel.innerHTML = exps.map(e => `<option value="${e}">${e}</option>`).join('');
            } else {
                sel.innerHTML = '<option value="">No expiries found</option>';
            }
        } catch (e) {
            sel.innerHTML = '<option value="">Error</option>';
        }
    },

    async inspectData() {
        const date = document.getElementById('selDate').value;
        const root = document.getElementById('selRoot').value;
        const expiry = document.getElementById('selExpiry').value;
        
        if (!date || !root || !expiry || root === 'Fetching...' || expiry === 'Fetching...') {
            alert("Please ensure Date, Symbol, and Expiry are selected properly.");
            return;
        }

        const btn = document.getElementById('btnInspect');
        btn.disabled = true;
        btn.innerHTML = "⏳ Fetching and Analyzing Chunks...";
        
        const payload = {
            date: date,
            root: root,
            expiry: expiry,
            atm_range: parseInt(document.getElementById('inpAtm').value),
            required_vars: document.getElementById('inpVars').value.split(',').map(s=>s.trim()),
            start_time: document.getElementById('inpStart').value,
            end_time: document.getElementById('inpEnd').value
        };

        try {
            const res = await fetch('/api/inspect_data', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(payload)
            });
            const data = await res.json();

            if(data.error) {
                alert("Inspection Error: " + data.error);
            } else {
                this.currentSession = data.session_id;
                this.currentMeta = data.meta;
                this.currentMeta.stats = data.stats;
                this.trimStart = data.stats.original_start;
                this.trimEnd = data.stats.original_end;
                this.originalStartStr = data.stats.original_start;
                this.originalEndStr = data.stats.original_end;
                this.missingTimeTicks = data.stats.missing_time_ticks;
                
                this.renderStats(data.stats);
                
                const c2 = document.getElementById('cardArchive');
                c2.style.opacity = '1';
                c2.style.pointerEvents = 'auto';
            }
        } catch(e) {
            alert(`Network Error: ${e.message}`);
        }
        btn.disabled = false;
        btn.innerHTML = "🔍 Fetch & Analyze Quality";
    },

    trimBlock(start, end, count, type, btnId) {
        if (type === 'start') {
            let endObj = new Date(end);
            endObj.setSeconds(endObj.getSeconds() + 10);
            this.trimStart = endObj.toISOString().replace('T',' ').substring(0, 19);
        } else if (type === 'end') {
            let startObj = new Date(start);
            startObj.setSeconds(startObj.getSeconds() - 10);
            this.trimEnd = startObj.toISOString().replace('T',' ').substring(0, 19);
        }
        
        this.missingTimeTicks -= count;
        document.getElementById('missingCountTxt').innerText = this.missingTimeTicks;
        document.getElementById('missingCountTxt').className = "stat-val " + (this.missingTimeTicks === 0 ? 'val-good' : 'val-warn');
        
        document.getElementById('bndStart').innerText = this.trimStart;
        document.getElementById('bndEnd').innerText = this.trimEnd;
        
        const b = document.getElementById(btnId);
        b.innerText = "Trimmed ✔";
        b.disabled = true;
        b.style.opacity = "0.5";
        b.style.pointerEvents = "none";
    },

    renderStats(s) {
        const mClass = this.missingTimeTicks === 0 ? 'val-good' : (this.missingTimeTicks < 50 ? 'val-warn' : 'val-bad');
        
        let nullHtml = '';
        for(let key in s.null_breakdown) {
            let v = s.null_breakdown[key];
            let c = v.count === 0 ? 'val-good' : 'val-bad';
            
            let details = v.count > 0 
                ? `<details><summary>Inspect ${v.count} Glitches</summary><div class="inspect-times">${v.times.join(', ')}${v.count>50?'...':''}</div></details>` 
                : `<div style="font-size:11px; color:var(--success); font-weight:600;">Perfect Data</div>`;
                
            nullHtml += `<div class="stat-item" style="flex-direction:column; align-items:flex-start; padding: 10px 0;">
                <div style="display:flex; justify-content:space-between; width:100%; margin-bottom: 6px;">
                    <span style="font-size:12px;">Glitched Nulls/Zeroes in <b>${key}</b></span>
                    <span class="stat-val ${c}">${v.count}</span>
                </div>
                ${details}
            </div>`;
        }

        let blocksHtml = '';
        if (s.missing_blocks.length > 0) {
            blocksHtml = `<div style="margin-top:16px; margin-bottom: 8px; font-weight:700; font-size:11px; color:var(--text-muted); text-transform:uppercase;">Missing Time Boundaries</div>`;
            s.missing_blocks.forEach((b, i) => {
                let btn = '';
                if (b.start === this.originalStartStr) {
                    btn = `<button id="tb_${i}" class="btn-warn" onclick="App.trimBlock('${b.start}', '${b.end}', ${b.count}, 'start', 'tb_${i}')">✂ Trim Start</button>`;
                } else if (b.end === this.originalEndStr) {
                    btn = `<button id="tb_${i}" class="btn-warn" onclick="App.trimBlock('${b.start}', '${b.end}', ${b.count}, 'end', 'tb_${i}')">✂ Trim End</button>`;
                }
                
                blocksHtml += `<div class="stat-item" style="font-size:11.5px; background: #fff; padding: 8px 12px; border-radius: 6px; border: 1px solid var(--border);">
                    <span style="color:var(--warn); font-weight:500;">${b.start.split(' ')[1]} ➔ ${b.end.split(' ')[1]} (${b.count} ticks)</span>
                    ${btn}
                </div>`;
            });
        }

        document.getElementById('statsBox').innerHTML = `
            <div class="stat-item"><span>Data Start Bound</span><span class="stat-val" id="bndStart" style="color:var(--accent);">${this.trimStart}</span></div>
            <div class="stat-item"><span>Data End Bound</span><span class="stat-val" id="bndEnd" style="color:var(--accent);">${this.trimEnd}</span></div>
            <div class="stat-item"><span>Raw Fetched Ticks</span><span class="stat-val">${s.found_ticks}</span></div>
            <div class="stat-item"><span>Missing Ticks (Internal)</span><span class="stat-val ${mClass}" id="missingCountTxt">${this.missingTimeTicks}</span></div>
            ${blocksHtml}
            <div style="margin-top:20px; margin-bottom: 4px; font-weight:700; font-size:11px; color:var(--text-muted); text-transform:uppercase;">Data Quality Breakdown</div>
            ${nullHtml}
        `;
    },

    async uploadData() {
        const btn = document.getElementById('btnUpload');
        const sessionName = document.getElementById('inpSessionName').value || "Default_Session";
        
        btn.disabled = true;
        btn.innerHTML = "☁️ Imputing & Uploading to Hugging Face...";

        try {
            const res = await fetch('/api/clean_and_upload', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    session_id: this.currentSession,
                    session_name: sessionName,
                    fill_method: document.getElementById('selFill').value,
                    trim_start: this.trimStart,
                    trim_end: this.trimEnd,
                    metadata: this.currentMeta
                })
            });
            const data = await res.json();

            if(data.error) {
                alert(data.error);
            } else {
                this.showToast(data.message);
                
                document.getElementById('cardArchive').style.opacity = '0.5';
                document.getElementById('cardArchive').style.pointerEvents = 'none';
                document.getElementById('statsBox').innerHTML = '<div style="text-align:center; color:var(--text-muted); font-style:italic; padding: 10px 0;">Run Inspection first...</div>';
                
                this.loadHFFolders();
            }
        } catch(e) {
            alert(`Error: ${e.message}`);
        }
        btn.disabled = false;
        btn.innerHTML = "☁️ Clean & Archive to Session";
    },

    async loadHFFolders() {
        const listDiv = document.getElementById('hfList');
        listDiv.innerHTML = '<div style="text-align:center; padding:40px; color:var(--text-muted);">Fetching repository structure...</div>';
        
        const res = await fetch('/api/hf_files');
        const data = await res.json();
        
        if(data.status === 'error') {
            listDiv.innerHTML = `<div style="color:var(--danger); text-align:center; padding:20px;">${data.message}</div>`;
            return;
        }

        const sessions = data.sessions;
        if(Object.keys(sessions).length === 0) {
            listDiv.innerHTML = '<div style="color:var(--text-muted); text-align:center; padding:40px;">No datasets archived yet. Create a session and archive data.</div>';
            return;
        }

        let html = '';
        for (const [sessionName, datasets] of Object.entries(sessions)) {
            html += `<div class="session-group">
                <div class="session-header">
                    <span>📁 ${sessionName} <span style="font-weight:400; font-size:12px; color:var(--text-muted);">(${datasets.length} folders)</span></span>
                    <div style="display:flex; gap:8px;">
                        <button class="btn-download-session" style="background:#0ea5e9;" onclick="App.downloadSessionParquetsZip('${sessionName}')">📊 DL Parquets ZIP</button>
                        <button class="btn-download-session" onclick="App.downloadSessionZip('${sessionName}')">📦 DL Full Session ZIP</button>
                    </div>
                </div>
                <div class="session-datasets">`;
                
            for (const ds of datasets) {
                html += `<div class="dataset-item">
                    <span class="dataset-name">📄 ${ds}</span>
                    <div class="hf-btns">
                        <button class="hf-btn" onclick="App.viewMeta('${sessionName}', '${ds}')">📝 Meta</button>
                        <button class="hf-btn" onclick="App.viewData('${sessionName}', '${ds}')">📊 Data Preview</button>
                    </div>
                </div>`;
            }
            html += `</div></div>`;
        }
        listDiv.innerHTML = html;
    },

    downloadSessionZip(sessionName) {
        this.showToast(`Zipping FULL session '${sessionName}'. Download will start shortly...`);
        window.location.href = `/api/download_session_zip?session_name=${encodeURIComponent(sessionName)}`;
    },
    
    downloadSessionParquetsZip(sessionName) {
        this.showToast(`Zipping PARQUETS ONLY for '${sessionName}'. Download will start shortly...`);
        window.location.href = `/api/download_session_parquets_zip?session_name=${encodeURIComponent(sessionName)}`;
    },

    async viewMeta(sessionName, datasetName) {
        document.getElementById('metaBox').style.display = 'block';
        document.getElementById('metaContent').innerText = 'Loading metadata...';
        
        const res = await fetch(`/api/hf_metadata?session_name=${encodeURIComponent(sessionName)}&dataset_name=${encodeURIComponent(datasetName)}`);
        const data = await res.json();
        
        if(data.status === 'error') {
            document.getElementById('metaContent').innerText = data.message;
        } else {
            document.getElementById('metaContent').innerText = JSON.stringify(data.metadata, null, 2);
        }
    },

    async viewData(sessionName, datasetName) {
        document.getElementById('dataModal').style.display = 'flex';
        document.getElementById('modalTitle').innerText = `${datasetName} (Preview: Max 1000 Rows)`;
        const body = document.getElementById('modalBody');
        body.innerHTML = '<div style="text-align:center; padding: 100px; color:var(--text-muted); font-size: 14px;">Downloading Parquet from Hugging Face & Generating Table...</div>';

        try {
            const res = await fetch(`/api/view_parquet?session_name=${encodeURIComponent(sessionName)}&dataset_name=${encodeURIComponent(datasetName)}`);
            const resp = await res.json();

            if (resp.status === 'error') {
                body.innerHTML = `<div style="color:var(--danger); text-align:center; padding:50px; font-weight:600;">${resp.message}</div>`;
                return;
            }

            let html = '<table class="data-table"><thead><tr>';
            resp.columns.forEach(c => html += `<th>${c}</th>`);
            html += '</tr></thead><tbody>';

            resp.data.forEach(row => {
                html += '<tr>';
                resp.columns.forEach(c => {
                    let val = row[c];
                    if(typeof val === 'number') val = val.toFixed(2);
                    html += `<td>${val}</td>`;
                });
                html += '</tr>';
            });
            html += '</tbody></table>';
            body.innerHTML = html;
        } catch (e) {
            body.innerHTML = `<div style="color:var(--danger); text-align:center; padding:50px;">Failed to load data preview.</div>`;
        }
    }
};

window.onload = () => App.init();
</script>
</body>
</html>
"""

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)