import http.server
import socketserver
import os
import json
import threading
import webbrowser
from pathlib import Path

# Relative to marco-package/marco/cli/web_serve.py
DIST_DIR = Path(__file__).parent.parent / "ui" / "dist"

def get_all_versions(repo_path: Path):
    versions_dir = repo_path / '.marco' / 'versions'
    if not versions_dir.exists():
        return {}
    
    result = {}
    for d in versions_dir.iterdir():
        if d.is_dir() and (d / 'manifest.json').exists():
            vid = d.name
            manifest = json.loads((d / 'manifest.json').read_text(encoding='utf-8'))
            metrics = json.loads((d / 'metrics.json').read_text(encoding='utf-8'))
            dag_structure = {}
            # The frontend expects 'metricsPerStep' explicitly
            metricsPerStep = metrics.get('metrics_per_step', {})

            # --- Extract Training Results ---
            training_results = None
            for p in d.iterdir():
                if p.name.startswith('training_results') or 'result' in p.name.lower():
                    training_results = p.read_text(encoding='utf-8', errors='ignore')
                    break

            # --- Inject Model Metrics Node ---
            from marco.model_drift import _parse_text_metrics
            try:
                # Determine final nodes in DAG to attach to
                depended_on = set()
                for k, v in dag_structure.items():
                    for dep in v.get('depends_on', []):
                        depended_on.add(dep)
                final_nodes = [k for k in dag_structure.keys() if k not in depended_on]
                
                # Use our robust regex parser directly on the text we just found
                model_metrics = _parse_text_metrics(training_results or "")
                
                # If parsed correctly, synthesize the node!
                if model_metrics:
                    acc = format(model_metrics.get('accuracy', 0.0), ".3f")
                    f1 = format(model_metrics.get('f1', 0.0), ".3f")
                    
                    dag_structure["model_evaluation"] = {
                        "func": "📊 Model Evaluation",
                        "depends_on": final_nodes,
                        "params": {
                            "Accuracy": f"{acc}",
                            "Macro F1": f"{f1}"
                        }
                    }
                    
                    metricsPerStep["model_evaluation"] = {
                        "Evaluation Report": training_results,
                        "Accuracy": acc,
                        "F1 Score": f1
                    }
            except Exception:
                pass
            # ---------------------------------

            # Raw pipeline stats (before processing)
            raw_stats = {}
            if (d / 'raw_stats.json').exists():
                try:
                    raw_stats = json.loads((d / 'raw_stats.json').read_text(encoding='utf-8'))
                except Exception:
                    pass

            # Post-pipeline stats (reductions + final state)
            post_stats = {}
            if (d / 'post_stats.json').exists():
                try:
                    post_stats = json.loads((d / 'post_stats.json').read_text(encoding='utf-8'))
                except Exception:
                    pass

            # Token frequency summary
            token_freq_summary = {}
            if (d / 'token_freq_summary.json').exists():
                try:
                    token_freq_summary = json.loads((d / 'token_freq_summary.json').read_text(encoding='utf-8'))
                except Exception:
                    pass
            
            # [Removed from here, hoisted up earlier]

            # Raw input data (first 200 chars per line to avoid huge payloads)
            raw_text = None
            if (d / 'raw.txt').exists():
                try:
                    lines = (d / 'raw.txt').read_text(encoding='utf-8', errors='ignore').splitlines()
                    raw_text = '\n'.join(lines[:60])  # cap at 60 rows
                except Exception:
                    pass

            # Preprocessed / cleaned output data
            preprocessed_data = None
            if (d / 'preprocessed.tsv').exists():
                try:
                    lines = (d / 'preprocessed.tsv').read_text(encoding='utf-8', errors='ignore').splitlines()
                    preprocessed_data = '\n'.join(lines[:60])  # cap at 60 rows
                except Exception:
                    pass

            # Pipeline config
            config_data = {}
            if (d / 'config.json').exists():
                try:
                    config_data = json.loads((d / 'config.json').read_text(encoding='utf-8'))
                except Exception:
                    pass
            
            result[vid] = {
                "manifest": manifest,
                "metrics": metrics,
                "metricsPerStep": metricsPerStep,
                "dagStructure": dag_structure,
                "trainingResults": training_results,
                "rawStats": raw_stats,
                "postStats": post_stats,
                "tokenFreqSummary": token_freq_summary,
                "rawText": raw_text,
                "preprocessedData": preprocessed_data,
                "config": config_data,
            }
    return result

def get_repo_meta(repo_path: Path):
    """Read all .marco repository-level metadata files."""
    marco = repo_path / '.marco'
    if not marco.exists():
        return None

    def load_json(p):
        try:
            return json.loads(p.read_text(encoding='utf-8')) if p.exists() else None
        except Exception:
            return None

    return {
        "refs":      load_json(marco / 'refs.json'),
        "chains":    load_json(marco / 'chains.json'),
        "registry":  load_json(marco / 'registry.json'),
        "root":      load_json(marco / 'root.json'),
        "workspace": load_json(marco / 'workspace.json'),
    }


class APIRequestHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/api/versions':
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header('Cache-Control', 'no-store, no-cache, must-revalidate')
            self.end_headers()
            
            try:
                data = get_all_versions(Path.cwd())
                response = json.dumps(data)
            except Exception as e:
                response = json.dumps({"error": str(e)})
                
            self.wfile.write(response.encode('utf-8'))
        else:
            return super().do_GET()
            
    def end_headers(self):
        self.send_header('Cache-Control', 'no-store, no-cache, must-revalidate')
        super().end_headers()

def serve(port: int = 7654, open_browser: bool = True):
    if not DIST_DIR.exists():
        print(f"Error: UI dist folder not found at {DIST_DIR}")
        print("Please ensure the react build is placed in marco/ui/dist")
        return

    user_cwd = Path.cwd()
    os.chdir(DIST_DIR)
    
    class ScopedHandler(APIRequestHandler):
        def do_GET(self):
            if self.path == '/api/versions':
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header('Cache-Control', 'no-store, no-cache, must-revalidate')
                self.end_headers()
                
                try:
                    data = get_all_versions(user_cwd)
                    response = json.dumps(data)
                except Exception as e:
                    response = json.dumps({"error": str(e)})
                    
                self.wfile.write(response.encode('utf-8'))
            elif self.path == '/api/repo':
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header('Cache-Control', 'no-store, no-cache, must-revalidate')
                self.end_headers()
                try:
                    meta = get_repo_meta(user_cwd)
                    response = json.dumps(meta or {})
                except Exception as e:
                    response = json.dumps({"error": str(e)})
                self.wfile.write(response.encode('utf-8'))
            else:
                super(APIRequestHandler, self).do_GET()
                
    with socketserver.TCPServer(("", port), ScopedHandler) as httpd:
        if open_browser:
            url = f"http://localhost:{port}"
            threading.Timer(0.5, lambda: webbrowser.open(url)).start()
        
        print(f"Marco React UI is running at URL http://localhost:{port}")
        print("Press Ctrl+C to stop.")
        try:
            # Run server in a separate daemon thread to allow main thread to catch KeyboardInterrupt
            server_thread = threading.Thread(target=httpd.serve_forever)
            server_thread.daemon = True
            server_thread.start()
            
            import time
            while True:
                time.sleep(0.5)
        except KeyboardInterrupt:
            print("\nShutting down React UI server.")
            httpd.shutdown()
            httpd.server_close()
            import sys
            sys.exit(0)

