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 = {}
            if (d / 'dag_structure.json').exists():
                dag_structure = json.loads((d / 'dag_structure.json').read_text(encoding='utf-8'))
            
            # The frontend expects 'metricsPerStep' explicitly
            metricsPerStep = metrics.get('metrics_per_step', {})
            
            result[vid] = {
                "manifest": manifest,
                "metrics": metrics,
                "metricsPerStep": metricsPerStep,
                "dagStructure": dag_structure
            }
    return result

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

    os.chdir(DIST_DIR)
    user_cwd = Path.cwd()
    
    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'))
            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:
            httpd.serve_forever()
        except KeyboardInterrupt:
            print("\nShutting down React UI server.")
