import argparse
import json
import os
import sys
from pathlib import Path
from typing import Optional

try:
    from rich.console import Console

    console = Console()
    HAS_RICH = True
except ImportError:
    HAS_RICH = False
    console = None


def print_msg(msg: str, style: str = None):
    if HAS_RICH and console:
        if style:
            console.print(f"[{style}]{msg}[/{style}]")
        else:
            console.print(msg)
    else:
        print(msg)


def main():
    parser = argparse.ArgumentParser(
        prog="verlex",
        description="Verlex - Run your code in the cloud for the price of a coffee",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  verlex login                    # Store API key
  verlex run script.py            # Run a script in the cloud
  verlex run script.py --gpu A100 # Run with specific GPU
  verlex jobs                     # List recent jobs
  verlex whoami                   # Show authenticated user

Get your API key at: https://verlex.dev
        """,
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    login_parser = subparsers.add_parser("login", help="Store API key")
    login_parser.add_argument("--api-key", "-k", help="API key")

    subparsers.add_parser("logout", help="Clear stored credentials")

    subparsers.add_parser("whoami", help="Show current authenticated user")

    run_parser = subparsers.add_parser("run", help="Run a script in the cloud")
    run_parser.add_argument("script", help="Python script to run")
    run_parser.add_argument(
        "--priority",
        action="store_true",
        default=True,
        help="Priority mode: immediate execution (default)",
    )
    run_parser.add_argument(
        "--patient",
        dest="priority",
        action="store_false",
        help="Patient mode: wait for lower price",
    )
    run_parser.add_argument("--gpu", "-g", help="GPU type (T4, A10, A100, H100, L4)")
    run_parser.add_argument("--cpu", "-c", type=int, help="CPU cores")
    run_parser.add_argument("--memory", "-m", help="Memory (e.g., 8GB, 16GB)")
    run_parser.add_argument("--timeout", "-t", type=int, default=3600, help="Timeout in seconds")

    jobs_parser = subparsers.add_parser("jobs", help="List recent jobs")
    jobs_parser.add_argument("--json", action="store_true", help="Output as JSON")

    status_parser = subparsers.add_parser("status", help="Check job status")
    status_parser.add_argument("job_id", help="Job ID")

    logs_parser = subparsers.add_parser("logs", help="View job logs")
    logs_parser.add_argument("job_id", help="Job ID")
    logs_parser.add_argument("--follow", "-f", action="store_true", help="Follow logs")

    cancel_parser = subparsers.add_parser("cancel", help="Cancel a running job")
    cancel_parser.add_argument("job_id", help="Job ID to cancel")

    subparsers.add_parser("version", help="Show version information")

    agent_parser = subparsers.add_parser(
        "agent",
        help="System agent: watch for heavy processes or submit scripts",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  verlex agent watch                   # Monitor system for heavy Python processes
  verlex agent watch --auto            # Auto-offload without prompting
  verlex agent watch --threshold 70    # Lower detection threshold to 70%
  verlex agent run script.py           # Submit a script via source-code pipeline
  verlex agent run train.py --gpu A100 # Submit with specific GPU
        """,
    )
    agent_sub = agent_parser.add_subparsers(dest="agent_command", help="Agent commands")

    watch_parser = agent_sub.add_parser(
        "watch", help="Watch system for heavy Python processes and offload to cloud"
    )
    watch_parser.add_argument(
        "--threshold", type=float, default=85.0,
        help="CPU/memory %% that triggers detection (default: 85)",
    )
    watch_parser.add_argument(
        "--interval", type=float, default=2.0,
        help="Seconds between scans (default: 2)",
    )
    watch_parser.add_argument(
        "--auto", action="store_true", default=False,
        help="Auto-offload without prompting",
    )
    watch_parser.add_argument(
        "--priority", action="store_true", default=False,
        help="Use priority mode (immediate execution)",
    )
    watch_parser.add_argument(
        "--flexible", action="store_true", default=True,
        help="Allow flexible resource allocation for lowest cost (default: True)",
    )
    watch_parser.add_argument("--gpu", "-g", help="GPU type for offloaded jobs")
    watch_parser.add_argument("--cpu", "-c", type=int, help="CPU cores for offloaded jobs")
    watch_parser.add_argument("--memory", "-m", help="Memory for offloaded jobs (e.g., 16GB)")

    agent_run_parser = agent_sub.add_parser(
        "run", help="Submit a script to the cloud via source-code pipeline"
    )
    agent_run_parser.add_argument("script", help="Python script to submit")
    agent_run_parser.add_argument(
        "--priority", action="store_true", default=False,
        help="Priority mode: immediate execution",
    )
    agent_run_parser.add_argument(
        "--patient", dest="priority", action="store_false",
        help="Patient mode: wait for lower price",
    )
    agent_run_parser.add_argument("--gpu", "-g", help="GPU type (T4, A10, A100, H100, L4)")
    agent_run_parser.add_argument("--cpu", "-c", type=int, help="CPU cores")
    agent_run_parser.add_argument("--memory", "-m", help="Memory (e.g., 8GB, 16GB)")
    agent_run_parser.add_argument(
        "--follow", "-f", action="store_true", default=True,
        help="Follow job logs until completion (default: True)",
    )
    agent_run_parser.add_argument(
        "--no-follow", dest="follow", action="store_false",
        help="Submit and return immediately",
    )

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        return 0

    try:
        if args.command == "login":
            return cmd_login(args)
        elif args.command == "logout":
            return cmd_logout(args)
        elif args.command == "whoami":
            return cmd_whoami(args)
        elif args.command == "run":
            return cmd_run(args)
        elif args.command == "jobs":
            return cmd_jobs(args)
        elif args.command == "status":
            return cmd_status(args)
        elif args.command == "logs":
            return cmd_logs(args)
        elif args.command == "cancel":
            return cmd_cancel(args)
        elif args.command == "version":
            return cmd_version(args)
        elif args.command == "agent":
            return cmd_agent(args, agent_parser)
        else:
            parser.print_help()
            return 1
    except KeyboardInterrupt:
        print("\nOperation cancelled.")
        return 130
    except Exception as e:
        print_msg(f"Error: {e}", "red")
        return 1


def get_api_key() -> Optional[str]:
    api_key = os.getenv("VERLEX_API_KEY")
    if api_key:
        return api_key

    config_path = Path.home() / ".verlex" / "config.json"
    if config_path.exists():
        try:
            with open(config_path) as f:
                config = json.load(f)
                return config.get("api_key")
        except Exception:
            pass

    return None


def save_api_key(api_key: str) -> None:
    config_dir = Path.home() / ".verlex"
    config_dir.mkdir(exist_ok=True)

    config_path = config_dir / "config.json"
    config = {}

    if config_path.exists():
        try:
            with open(config_path) as f:
                config = json.load(f)
        except Exception:
            pass

    config["api_key"] = api_key

    fd = os.open(config_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as f:
        json.dump(config, f, indent=2)


def cmd_login(args) -> int:
    import httpx

    api_key = args.api_key

    if not api_key:
        print("Enter your Verlex API key (get one at https://verlex.dev):")
        api_key = input("> ").strip()

    if not api_key:
        print_msg("No API key provided.", "red")
        return 1

    api_url = os.getenv("VERLEX_API_URL", "https://api.verlex.dev")

    try:
        response = httpx.get(
            f"{api_url}/v1/auth/me",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10.0,
        )

        if response.status_code == 200:
            save_api_key(api_key)
            data = response.json()
            print_msg(f"Logged in as: {data.get('email', 'user')}", "green")
            return 0
        elif response.status_code == 401:
            print_msg("Invalid API key.", "red")
            return 1
        else:
            print_msg(f"Login failed: {response.text}", "red")
            return 1
    except httpx.RequestError as e:
        print_msg(f"Network error: {e}", "red")
        return 1


def cmd_logout(args) -> int:
    config_path = Path.home() / ".verlex" / "config.json"

    if config_path.exists():
        try:
            with open(config_path) as f:
                config = json.load(f)
            config.pop("api_key", None)
            fd = os.open(config_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
            with os.fdopen(fd, "w") as f:
                json.dump(config, f, indent=2)
        except Exception as e:
            print_msg(f"Warning: Could not update config file: {e}", "yellow")

    print_msg("Logged out.", "green")
    return 0


def cmd_whoami(args) -> int:
    import httpx

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    api_url = os.getenv("VERLEX_API_URL", "https://api.verlex.dev")

    try:
        response = httpx.get(
            f"{api_url}/v1/auth/me",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10.0,
        )

        if response.status_code == 200:
            data = response.json()
            print(f"Email: {data.get('email', 'N/A')}")
            print(f"Tier:  {data.get('tier', 'free')}")
            print(f"Credits: ${data.get('credits', 0):.2f}")
            return 0
        else:
            print_msg("Failed to get user info.", "red")
            return 1
    except httpx.RequestError as e:
        print_msg(f"Network error: {e}", "red")
        return 1


def cmd_run(args) -> int:
    from verlex import GateWay, NotAuthenticatedError

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    script_path = Path(args.script)
    if not script_path.exists():
        print_msg(f"Script not found: {args.script}", "red")
        return 1

    with open(script_path) as f:
        script_code = f.read()

    def execute_script():
        exec(script_code, {"__name__": "__main__"})

    try:
        with GateWay(
            api_key=api_key,
            priority=args.priority,
            timeout=args.timeout,
        ) as gw:
            result = gw.run(
                execute_script,
                gpu=args.gpu,
                cpu=args.cpu,
                memory=args.memory,
            )

            if result is not None:
                print(f"\nResult: {result}")

            return 0

    except NotAuthenticatedError as e:
        print_msg("Invalid API key. Run 'verlex login' to re-authenticate.", "red")
        return 1
    except Exception as e:
        print_msg(f"Execution failed: {e}", "red")
        return 1


def cmd_jobs(args) -> int:
    import httpx

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    api_url = os.getenv("VERLEX_API_URL", "https://api.verlex.dev")

    try:
        response = httpx.get(
            f"{api_url}/v1/jobs",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10.0,
        )

        if response.status_code != 200:
            print_msg(f"Failed to get jobs: {response.text}", "red")
            return 1

        jobs = response.json().get("jobs", [])

        if args.json:
            print(json.dumps(jobs, indent=2))
            return 0

        if not jobs:
            print("No jobs found.")
            return 0

        print("\nRecent Jobs")
        print("=" * 60)
        print(f"{'JOB ID':<24} {'STATUS':<12} {'COST':<10}")
        print("-" * 60)

        for job in jobs[:20]:
            cost = job.get("cost", 0)
            cost_str = f"${cost:.4f}" if cost else "-"
            print(f"{job['id'][:22]:<24} {job['status']:<12} {cost_str:<10}")

        print()
        return 0

    except httpx.RequestError as e:
        print_msg(f"Network error: {e}", "red")
        return 1


def cmd_status(args) -> int:
    import httpx

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    api_url = os.getenv("VERLEX_API_URL", "https://api.verlex.dev")

    try:
        response = httpx.get(
            f"{api_url}/v1/jobs/{args.job_id}",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10.0,
        )

        if response.status_code == 404:
            print_msg(f"Job not found: {args.job_id}", "red")
            return 1
        elif response.status_code != 200:
            print_msg(f"Failed to get status: {response.text}", "red")
            return 1

        data = response.json()
        print(f"Job ID:   {args.job_id}")
        print(f"Status:   {data.get('status', 'unknown')}")
        if data.get("cost"):
            print(f"Cost:     ${data['cost']:.4f}")

        return 0

    except httpx.RequestError as e:
        print_msg(f"Network error: {e}", "red")
        return 1


def cmd_logs(args) -> int:
    import time

    import httpx

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    api_url = os.getenv("VERLEX_API_URL", "https://api.verlex.dev")

    def fetch_logs(offset: int = 0):
        response = httpx.get(
            f"{api_url}/v1/jobs/{args.job_id}/output",
            headers={"Authorization": f"Bearer {api_key}"},
            params={"offset": offset},
            timeout=10.0,
        )
        return response

    try:
        response = fetch_logs()

        if response.status_code == 404:
            print_msg(f"Job not found: {args.job_id}", "red")
            return 1
        elif response.status_code != 200:
            print_msg(f"Failed to get logs: {response.text}", "red")
            return 1

        data = response.json()
        entries = data.get("entries", [])

        for entry in entries:
            content = entry.get("content", entry.get("text", ""))
            print(content, end="")

        if args.follow:
            print("\n--- Following logs (Ctrl+C to stop) ---\n")
            last_index = data.get("current_index", 0)

            while True:
                time.sleep(1)
                response = fetch_logs(last_index)

                if response.status_code == 200:
                    data = response.json()
                    entries = data.get("entries", [])

                    for entry in entries:
                        content = entry.get("content", entry.get("text", ""))
                        print(content, end="")

                    if entries:
                        last_index = data.get("current_index", last_index)

                    if data.get("is_complete"):
                        print("\n--- Job completed ---")
                        break

        return 0

    except httpx.RequestError as e:
        print_msg(f"Network error: {e}", "red")
        return 1


def cmd_cancel(args) -> int:
    import httpx

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    api_url = os.getenv(
        "VERLEX_API_URL",
        "https://api.verlex.dev",
    )

    try:
        resp = httpx.post(
            f"{api_url}/v1/jobs/{args.job_id}/cancel",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10.0,
        )
        if resp.status_code == 200:
            print(f"Job {args.job_id} cancelled.")
            return 0
        elif resp.status_code == 400:
            print_msg(f"Job already finished.", "yellow")
            return 0
        elif resp.status_code == 404:
            print_msg(f"Job not found: {args.job_id}", "red")
            return 1
        else:
            print_msg(f"Cancel failed: {resp.text}", "red")
            return 1
    except httpx.RequestError as e:
        print_msg(f"Network error: {e}", "red")
        return 1


def cmd_version(args) -> int:
    from verlex import __version__

    print(f"Verlex CLI v{__version__}")
    print(f"Python {sys.version}")
    return 0


def cmd_agent(args, agent_parser) -> int:
    if not hasattr(args, "agent_command") or args.agent_command is None:
        agent_parser.print_help()
        return 0

    if args.agent_command == "watch":
        return cmd_agent_watch(args)
    elif args.agent_command == "run":
        return cmd_agent_run(args)
    else:
        agent_parser.print_help()
        return 1


def cmd_agent_watch(args) -> int:
    try:
        import psutil  # noqa: F401
    except ImportError:
        print_msg(
            "psutil is required for agent watch.  "
            "Install with:  pip install 'verlex[agent]'",
            "red",
        )
        return 1

    from verlex.agent import AgentDaemon

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    try:
        daemon = AgentDaemon(
            api_key=api_key,
            priority=args.priority,
            flexible=args.flexible,
            threshold=args.threshold,
            interval=args.interval,
            auto=args.auto,
            gpu=getattr(args, "gpu", None),
            cpu=getattr(args, "cpu", None),
            memory=getattr(args, "memory", None),
        )
        daemon.watch()
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print_msg(f"Agent watch failed: {e}", "red")
        return 1
    return 0


def cmd_agent_run(args) -> int:
    from verlex.agent import poll_job, submit_code

    api_key = get_api_key()
    if not api_key:
        print_msg("Not logged in. Run 'verlex login' first.", "yellow")
        return 1

    script_path = Path(args.script)
    if not script_path.exists():
        print_msg(f"Script not found: {args.script}", "red")
        return 1

    try:
        code = script_path.read_text(encoding="utf-8", errors="replace")
    except Exception as e:
        print_msg(f"Failed to read {args.script}: {e}", "red")
        return 1

    print(f"Submitting {script_path.name} to cloud (source-code pipeline)...")

    try:
        result = submit_code(
            code=code,
            file_path=str(script_path.resolve()),
            api_key=api_key,
            priority=args.priority,
            gpu=getattr(args, "gpu", None),
            cpu=getattr(args, "cpu", None),
            memory=getattr(args, "memory", None),
        )
    except Exception as e:
        print_msg(f"Submission failed: {e}", "red")
        return 1

    job_id = result.get("job_id", "unknown")
    print(f"  Job submitted: {job_id}")

    if not args.follow:
        print(f"  Track with:  verlex status {job_id}")
        print(f"  Logs:        verlex logs {job_id} --follow")
        return 0

    print(f"  Following logs...\n")

    try:
        final = poll_job(job_id=job_id, api_key=api_key, follow_logs=True)
    except KeyboardInterrupt:
        print(f"\n  Cancelling remote job...")
        try:
            import httpx

            api_url = os.environ.get(
                "VERLEX_API_URL",
                "https://api.verlex.dev",
            )
            resp = httpx.post(
                f"{api_url}/v1/jobs/{job_id}/cancel",
                headers={"Authorization": f"Bearer {api_key}"},
                timeout=10.0,
            )
            if resp.status_code == 200:
                print(f"  Job {job_id} cancelled.")
            else:
                print_msg(f"  Could not cancel job. Run: verlex cancel {job_id}", "red")
        except Exception:
            print_msg(f"  Could not cancel job. Run: verlex cancel {job_id}", "red")
        return 130
    except Exception as e:
        print_msg(f"\nFailed to follow job: {e}", "red")
        return 1

    status = final.get("status", "unknown")
    cost = final.get("cost") or final.get("actual_cost") or 0
    print(f"\n  Status: {status}")
    if cost:
        print(f"  Cost:   ${cost:.4f}")

    return 0 if status == "completed" else 1


if __name__ == "__main__":
    sys.exit(main())
