#!/usr/bin/env python3
"""
Helper script to extract Google Flow cookies and test the connection.

Usage:
  1. Open https://labs.google/fx/pt/tools/flow in Chrome
  2. Open DevTools (F12) → Console
  3. Paste this JavaScript and press Enter:

     copy(document.cookie)

  4. Run this script and paste the cookie string when prompted.

  OR manually copy from DevTools → Application → Cookies:
     - __Secure-next-auth.session-token
     - __Host-next-auth.csrf-token
"""

from __future__ import annotations

import os
import re
import sys
from pathlib import Path

ENV_PATH = Path(__file__).resolve().parent / ".env"


def extract_from_cookie_string(raw: str) -> dict[str, str]:
    """Parse a raw cookie string into the two tokens we need."""
    tokens = {}
    for part in raw.split(";"):
        part = part.strip()
        if part.startswith("__Secure-next-auth.session-token="):
            tokens["session"] = part.split("=", 1)[1]
        elif part.startswith("__Host-next-auth.csrf-token="):
            tokens["csrf"] = part.split("=", 1)[1]
    return tokens


def extract_from_json(raw: str) -> dict[str, str]:
    """Parse JSON cookie array (e.g. from EditThisCookie extension)."""
    import json as _json
    cookies = _json.loads(raw)
    tokens = {}
    for c in cookies:
        if c.get("name") == "__Secure-next-auth.session-token":
            tokens["session"] = c["value"]
        elif c.get("name") == "__Host-next-auth.csrf-token":
            tokens["csrf"] = c["value"]
    return tokens


def detect_and_extract(raw: str) -> dict[str, str]:
    """Auto-detect format (JSON array or cookie string) and extract."""
    raw = raw.strip()
    if raw.startswith("["):
        return extract_from_json(raw)
    return extract_from_cookie_string(raw)


def update_env(key: str, value: str) -> None:
    """Update a key in the .env file."""
    if not ENV_PATH.exists():
        ENV_PATH.write_text(f"{key}={value}\n", encoding="utf-8")
        return

    content = ENV_PATH.read_text(encoding="utf-8")
    if f"{key}=" in content:
        lines = content.splitlines()
        new_lines = []
        for line in lines:
            if line.startswith(f"{key}="):
                new_lines.append(f"{key}={value}")
            else:
                new_lines.append(line)
        ENV_PATH.write_text("\n".join(new_lines) + "\n", encoding="utf-8")
    else:
        with open(ENV_PATH, "a", encoding="utf-8") as f:
            f.write(f"\n{key}={value}\n")


def test_connection(session_token: str, csrf_token: str) -> bool:
    """Quick test: try to create a project to verify cookies work."""
    import httpx

    print("\n🔍 Testando conexão com Google Flow...")
    try:
        resp = httpx.get(
            "https://labs.google/fx/api/trpc/project.searchProjectWorkflows",
            params={
                "input": '{"json":{"pageSize":1,"projectId":"","toolName":"PINHOLE","fetchBookmarked":false,"rawQuery":"","mediaType":"MEDIA_TYPE_IMAGE","cursor":null},"meta":{"values":{"cursor":["undefined"]}}}'
            },
            cookies={
                "__Secure-next-auth.session-token": session_token,
                "__Host-next-auth.csrf-token": csrf_token,
            },
            timeout=15,
            follow_redirects=True,
        )
        if resp.status_code == 200:
            print("✅ Conexão OK! Cookies válidos.")
            return True
        elif resp.status_code == 401:
            print("❌ Cookies expirados ou inválidos (401 Unauthorized)")
            return False
        else:
            print(f"⚠️  Status inesperado: {resp.status_code}")
            print(f"   Response: {resp.text[:200]}")
            return False
    except Exception as e:
        print(f"❌ Erro de conexão: {e}")
        return False


def main():
    print("=" * 60)
    print("  Google Flow Cookie Extractor")
    print("=" * 60)
    print()
    print("Opção 1: Cole a string completa de cookies (document.cookie)")
    print("Opção 2: Cole os tokens individualmente")
    print()

    choice = input("Colar cookies (string ou JSON)? [S/n]: ").strip().lower()

    if choice != "n":
        print("\nCole os cookies (aceita JSON array ou cookie string):")
        print("(Para JSON multi-linha, cole tudo e pressione Enter duas vezes)")
        lines = []
        while True:
            line = input()
            if not line and lines:
                break
            lines.append(line)
        raw = "\n".join(lines).strip()
        tokens = detect_and_extract(raw)
        session_token = tokens.get("session", "")
        csrf_token = tokens.get("csrf", "")

        if not session_token:
            print("❌ __Secure-next-auth.session-token não encontrado na string!")
            print("   Tente a opção manual abaixo.")
            session_token = input("\nSession token: ").strip()
        if not csrf_token:
            print("⚠️  __Host-next-auth.csrf-token não encontrado.")
            print("   (Pode não estar em document.cookie por ser HttpOnly)")
            csrf_token = input("\nCSRF token (DevTools → Application → Cookies): ").strip()
    else:
        print("\nDevTools → Application → Cookies → https://labs.google")
        print()
        session_token = input("__Secure-next-auth.session-token: ").strip()
        csrf_token = input("__Host-next-auth.csrf-token: ").strip()

    if not session_token or not csrf_token:
        print("\n❌ Ambos os tokens são necessários!")
        sys.exit(1)

    # Test
    ok = test_connection(session_token, csrf_token)

    if ok:
        save = input("\n💾 Salvar no .env? [S/n]: ").strip().lower()
        if save != "n":
            update_env("GOOGLE_FLOW_SESSION_TOKEN", session_token)
            update_env("GOOGLE_FLOW_CSRF_TOKEN", csrf_token)
            print(f"✅ Tokens salvos em {ENV_PATH}")
    else:
        print("\n⚠️  Os cookies não funcionaram. Verifique se:")
        print("   1. Você está logado em https://labs.google/fx")
        print("   2. Os cookies foram copiados corretamente")
        print("   3. Os cookies não expiraram")

        force = input("\nSalvar mesmo assim? [s/N]: ").strip().lower()
        if force == "s":
            update_env("GOOGLE_FLOW_SESSION_TOKEN", session_token)
            update_env("GOOGLE_FLOW_CSRF_TOKEN", csrf_token)
            print(f"✅ Tokens salvos em {ENV_PATH}")


if __name__ == "__main__":
    main()