import os
import json
import re                             # регистронезависимый поиск email при VK-регистрации
import hashlib
import base64
from datetime import datetime, timedelta, UTC
import httpx
import logging
from bson import ObjectId
from config import *
from db import mongo_client, redis_client, gridfs_map, social_collection
import asyncio
import secrets
from utils import fetch_bytes, get_password_hash
import io
import uuid
from utils import encrypt, decrypt

# ---------------------- Константы OAuth ----------------------
OAUTH_URL = 'https://id.vk.com/oauth2/auth'
AUTHORIZE_URL = 'https://id.vk.ru/authorize'
SCOPES = 'offline,wall,photos,groups,email'   # email — для регистрации через VK
REDIRECT_URL = f'{FULL_DOMAIN}/api/v1/social/vk/token'


class TokenManager:
    """
    Управление VK-токенами с шифрованием чувствительных данных.
    """

    def __init__(self):
        self.tokens_coll = mongo_client["vk"]["tokens"]

    async def get_auth_url(self, user_id: str = None, purpose: str = "link", next: str = None) -> str:
        """purpose: "link" (привязка соцсети) или "register" (вход/регистрация через VK);
        next — куда вернуть пользователя после успешного входа."""
        state = os.urandom(32).hex()
        code_verifier = os.urandom(50).hex()
        code_challenge = base64.urlsafe_b64encode(
            hashlib.sha256(code_verifier.encode('utf-8')).digest()
        ).decode('utf-8').rstrip('=')

        oauth_data = {
            "user_id": user_id,
            "code_verifier": code_verifier,
            "purpose": purpose,
            "next": next,
        }
        await redis_client.setex(f"oauth_state:{state}", 60 * 15, json.dumps(oauth_data))

        url = (
            f"{AUTHORIZE_URL}?"
            f"response_type=code&client_id={CLIENT_ID}&scope={SCOPES}"
            f"&redirect_uri={REDIRECT_URL}&state={state}"
            f"&code_challenge={code_challenge}&code_challenge_method=S256"
        )
        return url

    async def save_vk(self, contents):
        grid = gridfs_map.get('social')
        filename = str(uuid.uuid4())
        stream = io.BytesIO(contents)
        return await grid.upload_from_stream(filename, stream)

    async def get_token(self, code: str, state: str, device_id: str = None) -> dict | None:
        """
        Принимает `code` (выданный VK после авторизации) и `state`
        (проверка, что запрос пришёл от того же клиента).
        Возвращает словарь токенов или `None` в случае ошибки.
        """
        # Проверка обязательных параметров
        if not state or not code:
            logging.error("[TokenManager] Missing code or state")
            return None

        # Читаем данные из Redis, ранее сохранённые в get_auth_url()
        key = f"oauth_state:{state}"
        oauth_data_raw = await redis_client.get(key)

        if not oauth_data_raw:
            logging.error(f"[TokenManager] State not found or expired: {state}")
            return None

        oauth_data = json.loads(oauth_data_raw)
        user_id = oauth_data.get("user_id")
        purpose = oauth_data.get("purpose", "link")
        next_url = oauth_data.get("next")
        created = False
        code_verifier = oauth_data["code_verifier"]

        # Выполняем запрос к VK для обмена кода на токены
        try:
            async with httpx.AsyncClient(timeout=15.0, verify=ctx, http2=True) as client:
                response = await client.post(
                    OAUTH_URL,
                    data={
                        'grant_type': 'authorization_code',
                        'code': code,
                        'code_verifier': code_verifier,
                        'client_id': CLIENT_ID,
                        'device_id': device_id or "",
                        'redirect_uri': REDIRECT_URL,
                        # Генерируем новый state только для логирования/отслеживания.
                        'state': os.urandom(32).hex(),
                    }
                )
                response.raise_for_status()
                token_data = response.json()
                access_token = token_data['access_token']

                reqp = await client.post(f"{VK_API_URI}users.get", data={
                    "access_token": access_token,
                    "fields": "photo_max_orig,screen_name",
                    "v": "5.199",
                })
                vk_profile = reqp.json()['response'][0]

                # Режим регистрации: находим или создаём пользователя ДО сохранения соцсетей,
                # чтобы дальше socials/tokens сохранились уже с реальным user_id.
                if purpose == "register" and not user_id:
                    user_id, created = await self._resolve_user_for_register(
                        vk_profile, token_data.get("email")
                    )

                reqg = await client.post(f"{VK_API_URI}groups.get", data={
                    'access_token':access_token,
                    'extended':1,
                    'filter':'admin,editor,moder',
                    'v':'5.199',
                })
                vk_groups = reqg.json()['response']['items']
                urls = [vk_profile['photo_max_orig']] + [i['photo_200'] for i in vk_groups]
                image_bytes = await asyncio.gather(*(fetch_bytes(client, u) for u in urls))
                # В режиме регистрации ставим аватар VK в профиль, если фото не задано вручную
                if purpose == "register" and user_id and image_bytes:
                    await self._maybe_set_vk_avatar(user_id, created, image_bytes[0])
                results = [await self.save_vk(i) for i in image_bytes]
                socials = [{
                    'name': f"{vk_profile['last_name']} {vk_profile['first_name']}",
                    'screen_name': vk_profile['screen_name'],
                    'id': vk_profile['id'],
                    'img': results[0],
                    'user_id': user_id,
                    'active': True,
                    'type': 'vk',
                }]
                for i,group in enumerate(vk_groups):
                    socials.append({
                        'name': group['name'],
                        'screen_name': group['screen_name'],
                        'id': group['id'],
                        'group': True,
                        'type': 'vk',
                        'active': False, # True if len(vk_groups) < 4 else False,
                        'img': results[i+1],
                        'user_id': user_id
                    }) 
                for doc in socials:
                    await mongo_client['social']['social'].update_one({'id': doc['id'],'user_id': doc['user_id']},
                        {'$set': doc},upsert=True)
                # Сохраняем полученные токены (Mongo + Redis)
                await self.save_tokens(user_id, token_data, device_id, client_id=CLIENT_ID)

                # Удаляем использованный state, чтобы его нельзя было переиспользовать.
                await redis_client.delete(key)

                return {
                    "user_id": user_id,
                    "registered": created,
                    "purpose": purpose,
                    "next": next_url,
                }

        # Обработка ошибок HTTP (напр., неверный код, истёк скоп)
        except httpx.HTTPStatusError as e:
            logging.error(f"[TokenManager] Token exchange failed: {e.response.text}")

        # Любые прочие исключения (сеть, тайм‑аут и т.д.)
        except Exception as e:
            logging.error(f"[TokenManager] Token exchange error: {e}")

        # На любой ошибке гарантируем удаление state (чистка)
        await redis_client.delete(key)
        return None


    async def _resolve_user_for_register(self, vk_profile: dict, email: str | None):
        """Находит пользователя по vk_id или email, иначе создаёт нового.

        Возвращает кортеж (user_id: str, created: bool).
        При совпадении по email — привязывает vk_id к найденному аккаунту.
        """
        users = mongo_client["auth"]["users"]
        vk_id = vk_profile.get("id")

        # 1) Уже привязанный VK-аккаунт
        existing = await social_collection.find_one({'id':vk_id},{'user_id':1})
        if existing:
            return existing["user_id"], False
        existing = await users.find_one({"vk_id": vk_id}, {"_id": 1})
        if existing:
            return str(existing["_id"]), False

        # 2) Существующий аккаунт с таким email — логиним и привязываем VK
        if email:
            existing = await users.find_one(
                {"email": {"$regex": f"^{re.escape(email)}$", "$options": "i"}},
                {"_id": 1},
            )
            if existing:
                await users.update_one({"_id": existing["_id"]}, {"$set": {"vk_id": vk_id}})
                return str(existing["_id"]), False

        # 3) Создаём нового пользователя (роль по умолчанию — художник)
        first = vk_profile.get("first_name", "") or ""
        last = vk_profile.get("last_name", "") or ""
        doc = {
            "username": f"{first} {last}".strip(),
            "female_name_lastname": f"{last} {first}".strip(),
            "email": email,
            "phone": None,
            "img": "/static/painter.svg",
            "roles": ["user"],
            "extra_scopes": [],
            "vk_id": vk_id,
            "client_id": "vk",
            "living_place": None,
            "offer_accepted": False,
            # Неиспользуемый случайный пароль: вход по email/паролю невозможен,
            # но login-эндпоинт не падает на None.startswith(...).
            "hashed_password": get_password_hash(secrets.token_urlsafe(32)),
        }
        res = await users.insert_one(doc)
        return str(res.inserted_id), True

    async def _maybe_set_vk_avatar(self, user_id: str, created: bool, avatar_bytes: bytes):
        """Ставит аватар из VK в профиль, если пользователь не задал своё фото.

        «Не задан вручную» = img пустой или дефолтная заглушка (/static/...).
        Для только что созданного пользователя ставим всегда.
        """
        if not avatar_bytes:
            return
        users = mongo_client["auth"]["users"]
        if not created:
            u = await users.find_one({"_id": ObjectId(user_id)}, {"img": 1})
            img = (u or {}).get("img")
            if img and not str(img).startswith("/static/"):
                return  # у пользователя уже своё фото — не трогаем
        try:
            grid = gridfs_map["users"]
            file_id = await grid.upload_from_stream(f"vk_{user_id}.jpg", io.BytesIO(avatar_bytes))
            await users.update_one({"_id": ObjectId(user_id)}, {"$set": {"img": str(file_id)}})
        except Exception as e:
            logging.error(f"[TokenManager] VK avatar save failed: {e}")

    # ====================== ШИФРОВАНИЕ ======================
    async def _encrypt_token_fields(self, token_data: dict) -> dict:
        """Шифрует access_token и refresh_token."""
        encrypted = token_data.copy()
        if "access_token" in encrypted and isinstance(encrypted["access_token"], (str, bytes)):
            at = encrypted["access_token"].encode("utf-8") if isinstance(encrypted["access_token"], str) else encrypted["access_token"]
            encrypted["encrypted_access_token"] = encrypt(at)
            encrypted.pop("access_token", None)
        if "refresh_token" in encrypted and isinstance(encrypted["refresh_token"], (str, bytes)):
            rt = encrypted["refresh_token"].encode("utf-8") if isinstance(encrypted["refresh_token"], str) else encrypted["refresh_token"]
            encrypted["encrypted_refresh_token"] = encrypt(rt)
            encrypted.pop("refresh_token", None)
        return encrypted

    async def _decrypt_token_fields(self, token_doc: dict) -> dict:
        """Расшифровывает токены."""
        decrypted = dict(token_doc)
        if "encrypted_access_token" in decrypted:
            try:
                decrypted["access_token"] = decrypt(decrypted["encrypted_access_token"]).decode("utf-8")
                decrypted.pop("encrypted_access_token", None)
            except Exception:
                pass
        if "encrypted_refresh_token" in decrypted:
            try:
                decrypted["refresh_token"] = decrypt(decrypted["encrypted_refresh_token"]).decode("utf-8")
                decrypted.pop("encrypted_refresh_token", None)
            except Exception:
                pass
        return decrypted

    async def save_tokens(self, user_id: str, token_data: dict, device_id: str = None, client_id: int = CLIENT_ID):
        expires_in = token_data.get("expires_in", 3600)
        expires_at = datetime.now(UTC) + timedelta(seconds=expires_in)

        encrypted_data = await self._encrypt_token_fields(token_data)

        doc = encrypted_data.copy()
        uid = ObjectId(user_id)
        doc.update({
            "user_id": uid,
            "expires_at": expires_at,
            "client_id": client_id or CLIENT_ID,
            "device_id": device_id,
        })

        token_json = json.dumps({
            **{k: v for k, v in doc.items() if k not in ["expires_at", "user_id"]},
            "expires_at": expires_at.isoformat(),
            "user_id": user_id,
        })

        await redis_client.setex(f"tokens:{user_id}", 60 * 60 * 24 * 60, token_json)

        # Зашифрованный access_token
        if "access_token" in token_data:
            at_bytes = token_data["access_token"].encode("utf-8") if isinstance(token_data["access_token"], str) else token_data["access_token"]
            enc_access = encrypt(at_bytes)
            await redis_client.setex(f"access_token:{user_id}", expires_in - 60, enc_access)

        doc.pop('user_id', None)
        await self.tokens_coll.update_one(
            {"user_id": uid},
            {"$set": doc, "$setOnInsert": {"user_id": [uid]}},
            upsert=True,
        )

    async def _load_from_mongo(self, user_id):
        doc = await self.tokens_coll.find_one(
            {"user_id": ObjectId(user_id)},
            {
                "expires_at": 1,
                "encrypted_access_token": 1,
                "encrypted_refresh_token": 1,
                "device_id": 1,
                "client_id": 1,
            },
        )
        if doc:
            return await self._decrypt_token_fields(doc)
        return None

    async def _cache_in_redis(self, user_id, token_doc):
        token_json = json.dumps({
            **{k: v for k, v in token_doc.items() if k not in ["expires_at", "user_id", "_id"]},
            "expires_at": token_doc['expires_at'].isoformat(),
            "user_id": user_id,
        })
        await redis_client.setex(f"tokens:{user_id}", 60 * 60 * 24 * 60, token_json)

    async def get_access_token(self, user_id: str, need_refresh: bool = False) -> str | None:
        # Короткий токен (зашифрованный)
        access_token_enc = await redis_client.get(f"access_token:{user_id}")
        if access_token_enc and not need_refresh:
            try:
                return decrypt(access_token_enc).decode("utf-8")
            except Exception:
                pass

        # Полный кэш
        cached = await redis_client.get(f"tokens:{user_id}")
        if cached:
            try:
                token_data = json.loads(cached)
                token_data = await self._decrypt_token_fields(token_data)
            except Exception:
                token_data = None
        else:
            token_data = None

        if isinstance(token_data, dict) and "error" in token_data:
            token_data = None

        if token_data is None:
            doc = await self._load_from_mongo(user_id)
            if not doc:
                return None
            await self._cache_in_redis(user_id, doc)
            token_data = doc

        # Проверка срока
        if isinstance(token_data.get("expires_at"), str):
            expires_at = datetime.fromisoformat(token_data["expires_at"].replace("Z", "+00:00"))
        else:
            expires_at = token_data.get("expires_at")

        if expires_at and expires_at.astimezone(UTC) > datetime.now(UTC) + timedelta(minutes=5) and not need_refresh:
            ttl = int((expires_at.astimezone(UTC) - datetime.now(UTC)).total_seconds()) - 30
            enc = encrypt(token_data["access_token"].encode("utf-8"))
            await redis_client.setex(f"access_token:{user_id}", ttl, enc)
            return token_data["access_token"]

        return await self.refresh_tokens(
            user_id,
            token_data.get("refresh_token"),
            token_data.get("device_id"),
            token_data.get('client_id', CLIENT_ID)
        )

    async def refresh_tokens(self, user_id: str, refresh_token: str, device_id: str, client_id: int = CLIENT_ID) -> str | None:
        if not refresh_token or not device_id:
            return None
        try:
            async with httpx.AsyncClient(timeout=15.0, verify=ctx, http2=True) as client:
                resp = await client.post(
                    OAUTH_URL,
                    data={
                        'grant_type': 'refresh_token',
                        'refresh_token': refresh_token,
                        'client_id': client_id or CLIENT_ID,
                        'device_id': device_id,
                        'state': os.urandom(32).hex(),
                    },
                )
                resp.raise_for_status()
                new_tokens = resp.json()

                # В случае ошибок VK вернёт поле `error`
                if 'error' in new_tokens:
                    logging.error(
                        f"[TokenManager] Refresh failed for user {user_id}: "
                        f"{new_tokens['error']} - {new_tokens.get('error_description', 'No description')}"
                    )
                    return None

                # Сохраняем fresh‑токены
                await self.save_tokens(user_id, new_tokens, device_id, client_id)
                return new_tokens.get("access_token")
        except Exception as e:
            logging.error(f"[TokenManager] Refresh failed: {e}")
            return None

    async def revoke_tokens(self, user_id: str):
        await self.tokens_coll.delete_one({"user_id": ObjectId(user_id)})
        await redis_client.delete(f"tokens:{user_id}", f"access_token:{user_id}")