import os
import json
import pandas as pd
from datetime import datetime, timezone

class ForensicAnalyzer:
    """Enhanced Backend Engine with Multi-Baseline Intelligence."""
    def __init__(self, session_path):
        self.session_path = session_path
        self.ui_log = os.path.join(session_path, "raw_ui.jsonl")
        self.net_log = os.path.join(session_path, "network_traffic.jsonl")
        self.ws_log = os.path.join(session_path, "websocket_traffic.jsonl")
        self.js_dir = os.path.join(session_path, "js_resources")
        self.stage2_dir = os.path.join(session_path, "stage2_processed")
        self.baselines = self._index_baselines()
        self.js_index = self._index_js_resources()
        
    def _parse_ts(self, ts_str):
        """Robustly parses ISO timestamps into aware UTC datetimes."""
        if not ts_str: return datetime.now(timezone.utc)
        try:
            # Handle 'Z' suffix and other variants
            ts_str = ts_str.replace("Z", "+00:00")
            dt = datetime.fromisoformat(ts_str)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt
        except:
            return datetime.now(timezone.utc)

    def _index_baselines(self):
        """Indexes all baseline files by their capture timestamp."""
        indexed = []
        files = [f for f in os.listdir(self.session_path) if f.startswith("baseline_") and f.endswith(".json")]
        for f in files:
            try:
                with open(os.path.join(self.session_path, f), 'r', encoding='utf-8') as b:
                    data = json.load(b)
                    indexed.append({
                        "file": f,
                        "ts": self._parse_ts(data.get("ts")),
                        "inventory": data.get("inventory", [])
                    })
            except: pass
        return sorted(indexed, key=lambda x: x["ts"])

    def _index_js_resources(self):
        """Builds a lookup index of strings (selectors/endpoints) found in JS files."""
        index = {}
        if not os.path.exists(self.js_dir): return index
        
        for f in os.listdir(self.js_dir):
            if not f.endswith(".js"): continue
            try:
                with open(os.path.join(self.js_dir, f), 'r', encoding='utf-8', errors='ignore') as js_file:
                    content = js_file.read()
                    # We store the file name for any significant string match
                    # This is a simple heuristic-based index
                    index[f] = content
            except: pass
        return index

    def _find_code_references(self, target_str):
        """Returns filenames of JS resources containing the target string."""
        if not target_str or len(target_str) < 3: return []
        refs = []
        # Clean selector for better matching (e.g. remove nth-child etc)
        clean_target = target_str.split(":")[0].split(" > ")[-1].replace("#", "").replace(".", "")
        
        for fname, content in self.js_index.items():
            if target_str in content or clean_target in content:
                refs.append(fname)
        return refs[:3] # Limit to top 3

    def _resolve_selector(self, selector, event_ts):
        """Finds the human-readable label for a CSS selector in the closest baseline."""
        best_b = None
        for b in self.baselines:
            if b["ts"] <= event_ts: best_b = b
            else: break
            
        if not best_b: return selector
        
        for item in best_b["inventory"]:
            if item.get("selector") == selector:
                # Prioritize label or placeholder over generic anchor_text
                td = item.get("text_data", {})
                label = td.get("label") or td.get("placeholder") or item.get("anchor_text")
                return f"'{label}' ({selector})"
        return selector

    def load_ui_events(self):
        events = []
        if os.path.exists(self.ui_log):
            with open(self.ui_log, 'r', encoding='utf-8') as f:
                for line in f:
                    try: 
                        evt = json.loads(line)
                        evt["dt"] = self._parse_ts(evt.get("system_ts"))
                        events.append(evt)
                    except: pass
        return pd.DataFrame(events)

    def load_network_traffic(self):
        traffic = []
        if os.path.exists(self.net_log):
            with open(self.net_log, 'r', encoding='utf-8') as f:
                for line in f:
                    try: traffic.append(json.loads(line))
                    except: pass
        return pd.DataFrame(traffic)

    def refine_actions(self, df_ui, df_net):
        """Converts raw UI events into a human-readable narrative with network context."""
        if df_ui.empty: return "No UI events recorded."
        
        # Ensure dt column exists
        if "dt" not in df_ui.columns:
            df_ui["dt"] = datetime.now(timezone.utc)
        
        narrative = []
        if "page_id" not in df_ui.columns:
            df_ui["page_id"] = "tab_unknown"

        for pid, group in df_ui.groupby("page_id"):
            narrative.append(f"\n--- Journey on {pid} ---")
            group = group.sort_values("dt", ascending=True)
            
            for _, row in group.iterrows():
                try:
                    ts = row["dt"].strftime("%H:%M:%S") if hasattr(row["dt"], "strftime") else "??:??:??"
                    type_ = row.get("type")
                    action = row.get("action", type_)
                    sel = row.get("sel", "")
                    val = row.get("val", "")
                    tid = row.get("trace_id")
                    
                    label = self._resolve_selector(sel, row["dt"]) if sel else ""
                    code_refs = self._find_code_references(sel) if sel else []

                    if action == "USER_CLICKED":
                        ref_str = f" [📜 Corel: {', '.join(code_refs)}]" if code_refs else ""
                        narrative.append(f"[{ts}] 🖱️ Clicked on {label}{ref_str}")
                    elif action == "USER_TYPING":
                        narrative.append(f"[{ts}] ⌨️ Typed '{val}' in {label}")
                    elif action == "USER_HOVER_OVER":
                        narrative.append(f"[{ts}] ✨ Hovered over {label}")
                    elif action == "USER_IDLE":
                        narrative.append(f"[{ts}] 😴 User went idle.")
                    elif action == "USER_ACTIVE":
                        narrative.append(f"[{ts}] ⏰ User returned after {row.get('was_idle_for_ms', 0)/1000:.1f}s")
                    elif type_ == "VALUE_DNA":
                        origin = row.get("origin")
                        narrative.append(f"[{ts}] 🧬 DNA: {label} updated to '{val}' by {origin}")
                    elif action == "TAB_OPENED":
                        narrative.append(f"[{ts}] 🌏 New Tab: {row.get('url')}")
                    elif type_ == "POPUP_SHOW":
                        narrative.append(f"[{ts}] 🎯 Popup Detected: {row.get('text', '')}")

                    # Inline Network Correlation via Trace ID
                    if tid and not df_net.empty:
                        related_net = df_net[df_net["trace_id"] == tid]
                        for _, nreq in related_net.iterrows():
                            url = nreq.get("url", "")
                            # Simplify URL for readability
                            simplified_url = url.split("?")[0][-50:]
                            
                            # Find JS that might have triggered this URL
                            endpoint = url.split("?")[0].split("/")[-1]
                            net_code_refs = self._find_code_references(endpoint) if len(endpoint) > 3 else []
                            net_ref_str = f" [📜 Trigger: {', '.join(net_code_refs)}]" if net_code_refs else ""
                            
                            narrative.append(f"      ↳ 🌐 Network: {nreq.get('method')} ...{simplified_url}{net_ref_str}")
                except: continue

        return "\n".join(narrative)

    def generate_report(self):
        """Produces a comprehensive forensic report."""
        df_ui = self.load_ui_events()
        df_net = self.load_network_traffic()
        
        report = {
            "session": os.path.basename(self.session_path),
            "timestamp": datetime.now().isoformat(),
            "summary": {
                "ui_events": len(df_ui),
                "network_calls": len(df_net),
                "baselines_captured": len(self.baselines)
            },
            "trust_audit": self._run_trust_audit(df_ui),
            "narrative": self.refine_actions(df_ui, df_net)
        }
        
        os.makedirs(self.stage2_dir, exist_ok=True)
        report_path = os.path.join(self.stage2_dir, "analysis_report.json")
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False)
        
        print(f"\n✅ Analysis Complete: {report_path}")
        print(f"🧬 Trust Score: {report['trust_audit'].get('trust_score', 'N/A')}")
        print("-" * 50)
        print(report["narrative"][:2000] + ("..." if len(report["narrative"]) > 2000 else ""))

    def _run_trust_audit(self, df_ui):
        if df_ui.empty: return {}
        dna = df_ui[df_ui["type"] == "VALUE_DNA"]
        if dna.empty: return {"status": "NO_DNA_DATA"}
        stats = dna["origin"].value_counts().to_dict()
        trust_score = (stats.get("USER_TYPED", 0) / dna.shape[0]) * 100 if dna.shape[0] > 0 else 0
        return {"total_mutations": int(dna.shape[0]), "origins": stats, "trust_score": f"{trust_score:.1f}%"}

def main():
    session_root = "sessions"
    if not os.path.exists(session_root): return
    sessions = sorted([(os.path.getmtime(os.path.join(session_root, d)), os.path.join(session_root, d)) 
                       for d in os.listdir(session_root) if os.path.isdir(os.path.join(session_root, d))])
    if sessions:
        latest = sessions[-1][1]
        print(f"🧐 Analyzing latest session: {latest}")
        analyzer = ForensicAnalyzer(latest)
        analyzer.generate_report()

if __name__ == "__main__":
    main()