# orchestrator_text.py

import subprocess
import sys
from pathlib import Path
import argparse
import shutil
import json
from typing import Optional, Dict, List, Union

# --- Configuration ---
PIPELINE_SCRIPTS = [
    "3_design_generator.py", 
    "4_code_generator.py"
]

BASE_OUTPUT_DIR = "../data/website_data"

# Available models
AVAILABLE_MODELS = {
    0: "Gemini 2.5 Pro (GA)",
    1: "Gemini 2.5 Flash (GA)", 
    2: "Gemini 2.5 Flash-Lite (Preview)"
}

# --- MODEL CONFIGURATION ---
MODEL_CONFIG = {
    "design": 0,
    "code": 0
}

def run_step_with_live_output(command: List[str]) -> Optional[Dict[str, Union[float, int]]]:
    """
    Runs a script as a subprocess and captures the final JSON line for cost analysis.
    """
    try:
        import os
        env = os.environ.copy()
        env['PYTHONUNBUFFERED'] = '1'
        
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding='utf-8',
            bufsize=1,
            universal_newlines=True,
            env=env
        )

        last_line = ""
        if process.stdout:
            for line in iter(process.stdout.readline, ''):
                print(line, end='', flush=True)
                last_line = line.strip() if line.strip() else last_line

        process.wait()

        if process.returncode != 0:
            print(f"\n❌ ERROR: Script '{Path(command[1]).name}' failed with exit code {process.returncode}.")
            return None

        try:
            return json.loads(last_line) if last_line else None
        except json.JSONDecodeError:
            return None

    except FileNotFoundError:
        print(f"❌ Error: Could not find script at '{command[1]}'")
        return None
    except Exception as e:
        print(f"❌ An unexpected error occurred: {e}")
        return None

def get_script_command(script_name: str, site_dir: str) -> List[str]:
    """Generate the appropriate command for each script."""
    base_command = [sys.executable, script_name, site_dir]
    if script_name == "3_design_generator.py":
        return base_command + ["--model-index", str(MODEL_CONFIG['design'])]
    elif script_name == "4_code_generator.py":
        return base_command + ["--model-index", str(MODEL_CONFIG['code'])]
    return base_command

def main():
    parser = argparse.ArgumentParser(description="Runs the text-based website generation pipeline.")
    parser.add_argument("--domain", required=True, type=str, help="The domain name for the new site (e.g., company.com).")
    parser.add_argument("--company_name", required=True, type=str, help="The name of the company.")
    parser.add_argument("--site_brief", required=True, type=str, help="A detailed brief about the company and design requirements.")
    
    args = parser.parse_args()

    site_output_dir = Path(BASE_OUTPUT_DIR) / args.domain

    # Create the site directory and initial files
    print(f"🚀 Creating new site for: {args.company_name} ({args.domain})")
    site_output_dir.mkdir(parents=True, exist_ok=True)
    
    # Create a basic manifest.json
    manifest = {"domain": args.domain, "company_name": args.company_name}
    with open(site_output_dir / "manifest.json", "w") as f:
        json.dump(manifest, f, indent=2)

    # Create the site_brief.json from the input with proper structure expected by design generator
    site_brief = {
        "company_name": args.company_name,
        "domain": args.domain,
        "content_brief": args.site_brief,
        "brand_profile": {
            "brand_identity": args.company_name,
            "business_keywords": ["business", "professional", "service"],
            "colors": [
                {"name": "Primary Blue", "hex": "#007bff"},
                {"name": "Secondary Gray", "hex": "#6c757d"},
                {"name": "Success Green", "hex": "#28a745"},
                {"name": "Dark Text", "hex": "#333333"}
            ],
            "typography": {
                "primary_font": "Inter",
                "secondary_font": "system-ui",
                "styles": ["clean", "modern", "readable"]
            },
            "design_style": "Clean, modern, and professional design suitable for business websites",
            "logo_description": f"Modern text-based logo for {args.company_name}",
            "visual_elements": ["clean lines", "professional layout", "modern typography"]
        },
        "navigation_links": [
            {"text": "Home", "href": "#home"},
            {"text": "About", "href": "#about"},
            {"text": "Services", "href": "#services"},
            {"text": "Contact", "href": "#contact"}
        ],
        "site_metadata": {
            "site_title": args.company_name,
            "primary_domain": args.domain,
            "meta_description": f"Professional website for {args.company_name}"
        },
        "site_pages": [
            {
                "page_key": "home",
                "title": "Home",
                "url": "/",
                "content_sections": {
                    "hero": f"Welcome to {args.company_name}",
                    "about": args.site_brief,
                    "services": "Our professional services and offerings",
                    "contact": "Get in touch with us today"
                }
            },
            {
                "page_key": "about",
                "title": "About Us",
                "url": "/about",
                "content_sections": {
                    "main_content": args.site_brief
                }
            },
            {
                "page_key": "contact",
                "title": "Contact",
                "url": "/contact",
                "content_sections": {
                    "contact_info": "Contact information and details"
                }
            }
        ]
    }
    with open(site_output_dir / "site_brief.json", "w") as f:
        json.dump(site_brief, f, indent=2)
        
    print(f"✅ Created manifest and site brief for {args.domain}")

    total_pipeline_cost = 0.0

    print(f"🤖 Model Configuration:")
    print(f"   - Design Generation: {AVAILABLE_MODELS[MODEL_CONFIG['design']]}")
    print(f"   - Code Generation: {AVAILABLE_MODELS[MODEL_CONFIG['code']]}")
    print()

    for step_num, script_name in enumerate(PIPELINE_SCRIPTS, 1):
        print(f"\n--- STEP {step_num} of {len(PIPELINE_SCRIPTS)}: {script_name} ---")
        
        command = get_script_command(script_name, str(site_output_dir))
        usage = run_step_with_live_output(command)
        
        if usage is None:
            print(f"   Stopping pipeline for {args.domain} due to error in {script_name}")
            break
        
        script_cost = usage.get("cost", 0)
        total_pipeline_cost += script_cost
        
        if script_cost > 0:
            print(f"   💰 Step cost: ${script_cost:.6f}")

    else:
        print(f"\n✔️ Pipeline completed successfully for: {args.domain}")
    
    print(f"\n💰 Total Estimated Cost: ${total_pipeline_cost:.6f}")

if __name__ == "__main__":
    main()