import json
import logging
import requests
import yaml
import os
import time
import random

logger = logging.getLogger(__name__)


class PromptDesigner:
    def __init__(self, omniroute_url, omniroute_key, model="cohere/command-r-plus-08-2024",
                 timeout=30, temperature=0.7):
        self.url = omniroute_url.rstrip("/") + "/chat/completions"
        self.key = omniroute_key
        self.model = model
        self.timeout = timeout
        self.temperature = temperature
        self.headers = {
            "Authorization": f"Bearer {self.key}",
            "Content-Type": "application/json"
        }

    def _call_llm(self, system_prompt, user_prompt, max_tokens=500):
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            "max_tokens": max_tokens,
            "temperature": self.temperature
        }

        try:
            resp = requests.post(self.url, headers=self.headers, json=payload, timeout=self.timeout)
            resp.raise_for_status()
            data = resp.json()
            content = data["choices"][0]["message"]["content"].strip()
            tokens = data.get("usage", {}).get("total_tokens", 0)
            logger.info(f"Command R+ respondio: {len(content)} chars, {tokens} tokens")
            return content, tokens
        except requests.exceptions.Timeout:
            logger.error("Timeout llamando a Command R+")
            raise Exception("Command R+ no responde (timeout)")
        except Exception as e:
            logger.error(f"Error llamando a Command R+: {e}")
            raise

    def design_prompt(self, character_name, character_tags, spice_level="suggestive",
                      outfit_hint=None, scene_hint=None, extra_context=""):
        system_prompt = (
            "You are an expert Stable Diffusion prompt engineer specializing in anime NSFW content. "
            "You create detailed prompts using Danbooru tag format. "
            "Output ONLY comma-separated tags, no explanations, no numbering, no quotes. "
            "Structure: [quality tags], [character], [outfit], [pose], [details], [scenario]. "
            "Always include: masterpiece, best quality, highly detailed. "
            "Use the character tags provided EXACTLY - do not modify or invent new character tags. "
            "Be creative with poses, outfits, and scenarios while keeping the character accurate."
        )

        spice_desc = "suggestive with artisticcleavage, not explicit" if spice_level == "suggestive" else "explicit NSFW, undressing, bare skin"

        user_parts = [f"Character: {character_name}"]
        user_parts.append(f"Character tags (use EXACTLY): {character_tags}")
        user_parts.append(f"Spice level: {spice_desc}")

        if outfit_hint:
            user_parts.append(f"Outfit idea: {outfit_hint}")
        if scene_hint:
            user_parts.append(f"Scene idea: {scene_hint}")
        if extra_context:
            user_parts.append(f"Extra context: {extra_context}")

        user_parts.append("")
        user_parts.append("Generate ONE detailed Danbooru-style prompt. Output ONLY the comma-separated tags.")

        user_prompt = "\n".join(user_parts)

        content, tokens = self._call_llm(system_prompt, user_prompt, max_tokens=300)
        return content, tokens

    def design_batch(self, character_name, character_tags, spice_level="suggestive",
                     count=3, outfit_hint=None, scene_hint=None):
        system_prompt = (
            "You are an expert Stable Diffusion prompt engineer specializing in anime NSFW content. "
            "You create detailed prompts using Danbooru tag format. "
            "Each prompt must be on a separate line. Output ONLY comma-separated tags per line. "
            "No numbering, no explanations, no quotes. "
            "Structure: [quality tags], [character], [outfit], [pose], [details], [scenario]. "
            "Always include: masterpiece, best quality, highly detailed. "
            "Use the character tags provided EXACTLY. "
            "Be creative and vary poses, outfits, lighting, and scenarios across the prompts."
        )

        spice_desc = "suggestive with cleavage, artistic" if spice_level == "suggestive" else "explicit NSFW, undressing"

        user_parts = [f"Character: {character_name}"]
        user_parts.append(f"Character tags (use EXACTLY): {character_tags}")
        user_parts.append(f"Spice level: {spice_desc}")
        user_parts.append(f"Generate {count} DIFFERENT prompts, each on a separate line.")
        user_parts.append("")
        user_parts.append("Output ONLY the prompts, one per line:")

        user_prompt = "\n".join(user_parts)

        content, tokens = self._call_llm(system_prompt, user_prompt, max_tokens=count * 150)

        lines = [l.strip() for l in content.strip().split("\n") if l.strip() and not l.strip().startswith(("#", "1.", "2.", "3.", "4.", "5."))]
        return lines[:count], tokens

    def check_connection(self):
        try:
            content, tokens = self._call_llm(
                "You are a helpful assistant.",
                "Reply with just: OK",
                max_tokens=10
            )
            return True, content
        except Exception as e:
            return False, str(e)


def load_omniroute_config(base_dir):
    config_path = os.path.join(base_dir, "Config", "plumloom_engine.yaml")
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    omni_cfg = config["endpoints"]["omniroute"]
    api_key = os.environ.get("OMNIROUTE_API_KEY", "sk-9e3027503ec29237-25cded-bd502dd0")

    return {
        "url": omni_cfg["url"],
        "key": api_key,
        "timeout": omni_cfg.get("timeout", 30)
    }


class LiterouterDesigner:
    def __init__(self, api_keys, models=None, timeout=30, temperature=0.7):
        self.base_url = "https://api.literouter.com/v1/chat/completions"
        self.api_keys = api_keys if isinstance(api_keys, list) else [api_keys]
        self.models = models or [
            "deepseek-v3.2:free",
            "deepseek-v4-flash:free",
            "glm-5.3-flash:free",
            "qwen3.5:free"
        ]
        self.timeout = timeout
        self.temperature = temperature
        self.current_key_index = 0
        self.current_model_index = 0
        self.last_request_time = 0
        self.min_request_interval = 35  # seconds between requests

    def _get_next_key(self):
        key = self.api_keys[self.current_key_index]
        self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
        return key

    def _get_current_model(self):
        return self.models[self.current_model_index]

    def _rotate_model(self):
        self.current_model_index = (self.current_model_index + 1) % len(self.models)
        logger.info(f"Literouter: modelo rotado a {self._get_current_model()}")

    def _wait_if_needed(self):
        elapsed = time.time() - self.last_request_time
        if elapsed < self.min_request_interval:
            wait_time = self.min_request_interval - elapsed
            logger.info(f"Literouter: esperando {wait_time:.1f}s por rate limit")
            time.sleep(wait_time)

    def _call_llm(self, system_prompt, user_prompt, max_tokens=500):
        self._wait_if_needed()

        key = self._get_next_key()
        model = self._get_current_model()

        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            "max_tokens": max_tokens,
            "temperature": self.temperature
        }

        headers = {
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json"
        }

        try:
            resp = requests.post(self.base_url, headers=headers, json=payload, timeout=self.timeout)
            self.last_request_time = time.time()

            if resp.status_code == 429:
                logger.warning(f"Literouter: rate limit en modelo {model}, rotando...")
                self._rotate_model()
                return self._call_llm(system_prompt, user_prompt, max_tokens)

            resp.raise_for_status()
            data = resp.json()
            content = data["choices"][0]["message"]["content"].strip()
            tokens = data.get("usage", {}).get("total_tokens", 0)
            logger.info(f"Literouter ({model}) respondio: {len(content)} chars, {tokens} tokens")
            return content, tokens

        except requests.exceptions.Timeout:
            logger.error(f"Literouter: timeout con modelo {model}")
            self._rotate_model()
            raise Exception(f"Literouter timeout ({model})")
        except Exception as e:
            logger.error(f"Literouter error ({model}): {e}")
            if "402" in str(e) or "401" in str(e):
                logger.info("Literouter: modelo no disponible, rotando...")
                self._rotate_model()
            raise

    def design_prompt(self, character_name, character_tags, spice_level="suggestive",
                      outfit_hint=None, scene_hint=None, extra_context=""):
        system_prompt = (
            "You are an expert Stable Diffusion prompt engineer for anime NSFW content. "
            "Generate prompts in Danbooru tag format. "
            "CRITICAL RULES:\n"
            "1. ALL tags MUST be in ENGLISH ONLY - NEVER use Chinese, Japanese, or any other language\n"
            "2. Use underscores for multi-word tags (e.g. 'long_hair' not 'long hair')\n"
            "3. Output ONLY comma-separated tags, no explanations, no numbering, no quotes\n"
            "4. Use the character tags provided EXACTLY - do not modify character tags\n"
            "5. Tags must follow this EXACT structure with BREAK separators:\n"
            "   [character_tags], [outfit], [details], [lighting]\n"
            "   BREAK\n"
            "   [pose], [action], [expression]\n"
            "   BREAK\n"
            "   [scenario], [background]\n"
            "6. Pose tags are CRITICAL - put them RIGHT AFTER the first BREAK\n"
            "7. Be creative with poses, outfits and scenarios while keeping character accuracy\n"
            "8. NEVER include mirror, reflection, mirrored, glass_reflection, reflective_surface in prompts"
        )

        spice_desc = "suggestive with artistic cleavage, not explicit" if spice_level == "suggestive" else "explicit NSFW, undressing, bare skin"

        user_parts = [f"Character: {character_name}"]
        user_parts.append(f"Character tags (use EXACTLY): {character_tags}")
        user_parts.append(f"Spice level: {spice_desc}")

        if outfit_hint:
            user_parts.append(f"Outfit idea: {outfit_hint}")
        if scene_hint:
            user_parts.append(f"Scene idea: {scene_hint}")
        if extra_context:
            user_parts.append(f"Extra context: {extra_context}")

        user_parts.append("")
        user_parts.append("Generate ONE prompt with BREAK separators as specified. ALL tags in ENGLISH ONLY.")

        user_prompt = "\n".join(user_parts)

        content, tokens = self._call_llm(system_prompt, user_prompt, max_tokens=400)
        return content, tokens

    def design_batch(self, character_name, character_tags, spice_level="suggestive",
                     count=3, outfit_hint=None, scene_hint=None):
        system_prompt = (
            "You are an expert Stable Diffusion prompt engineer for anime NSFW content. "
            "Generate prompts in Danbooru tag format. "
            "CRITICAL RULES:\n"
            "1. ALL tags MUST be in ENGLISH ONLY - NEVER use Chinese, Japanese, or any other language\n"
            "2. Use underscores for multi-word tags (e.g. 'long_hair' not 'long hair')\n"
            "3. Each prompt on a separate line. Output ONLY comma-separated tags per line\n"
            "4. No numbering, no explanations, no quotes\n"
            "5. Use the character tags provided EXACTLY\n"
            "6. Each prompt must follow this structure with BREAK separators:\n"
            "   [character_tags], [outfit], [details], [lighting]\n"
            "   BREAK\n"
            "   [pose], [action], [expression]\n"
            "   BREAK\n"
            "   [scenario], [background]\n"
            "7. Pose tags are CRITICAL - put them RIGHT AFTER the first BREAK\n"
            "8. Be creative and vary poses, outfits, lighting, and scenarios across prompts\n"
            "9. NEVER include mirror, reflection, mirrored, glass_reflection, reflective_surface"
        )

        spice_desc = "suggestive with cleavage" if spice_level == "suggestive" else "explicit NSFW, undressing"

        user_parts = [f"Character: {character_name}"]
        user_parts.append(f"Character tags (use EXACTLY): {character_tags}")
        user_parts.append(f"Spice level: {spice_desc}")
        user_parts.append(f"Generate {count} DIFFERENT prompts, each on a separate line.")
        user_parts.append("")
        user_parts.append("Output ONLY prompts with BREAK separators. ALL tags in ENGLISH ONLY:")

        user_prompt = "\n".join(user_parts)

        content, tokens = self._call_llm(system_prompt, user_prompt, max_tokens=count * 200)

        lines = [l.strip() for l in content.strip().split("\n") if l.strip() and not l.strip().startswith(("#", "1.", "2.", "3.", "4.", "5."))]
        return lines[:count], tokens

    def design_progression(self, character_name, character_tags, outfit_tags, scene_tags=None):
        system_prompt = (
            "You are an expert Stable Diffusion prompt engineer for anime NSFW content. "
            "Generate a 4-step PROGRESSION SEQUENCE for a character going from fully dressed to NSFW. "
            "CRITICAL RULES:\n"
            "1. ALL tags MUST be in ENGLISH ONLY\n"
            "2. Use underscores for multi-word tags\n"
            "3. Each prompt on a separate line (4 lines total)\n"
            "4. Output ONLY comma-separated tags per line, no numbering, no explanations\n"
            "5. Use the character tags provided EXACTLY\n"
            "6. Each prompt must follow BREAK structure:\n"
            "   [character_tags], [outfit], [details], [lighting]\n"
            "   BREAK\n"
            "   [pose], [action], [expression]\n"
            "   BREAK\n"
            "   [scenario], [background]\n"
            "7. NEVER include mirror, reflection, mirrored, glass_reflection\n"
            "\n"
            "PROGRESSION STEPS (each line = one step):\n"
            "LINE 1 - OUTFIT: Character in FULL outfit, posing normally. Include full outfit tags (school uniform, dress, etc).\n"
            "LINE 2 - UNDRESSING: Character partially undressing. Use tags like: clothes_lift, shirt_lift, skirt_lift, unbuttoned, zipper_partial, panties_visible, bra_visible, thighhighs, stockings, one_breast_exposed.\n"
            "LINE 3 - NUDE/UNDERWEAR: Character in underwear or mostly nude. Use tags like: bikini, lingerie, bra, panties, topless, bottomless, nude, naked, sheer_clothes, see-through.\n"
            "LINE 4 - NSFW: Explicit sexual scene. Use tags like: sex, nude, naked, fucking, blowjob, cum, creampie, orgasm, moaning, penetration.\n"
            "\n"
            "IMPORTANT: Keep SAME pose/scenario across all 4 steps for visual continuity. "
            "The progression should feel natural - as if watching a sequence."
        )

        user_parts = [f"Character: {character_name}"]
        user_parts.append(f"Character tags (use EXACTLY): {character_tags}")
        user_parts.append(f"Base outfit: {outfit_tags}")
        if scene_tags:
            user_parts.append(f"Scene: {scene_tags}")
        user_parts.append("")
        user_parts.append("Generate the 4-step progression. Output ONLY the 4 prompts, one per line:")

        user_prompt = "\n".join(user_parts)

        content, tokens = self._call_llm(system_prompt, user_prompt, max_tokens=800)

        lines = [l.strip() for l in content.strip().split("\n") if l.strip() and not l.strip().startswith(("#", "1.", "2.", "3.", "4.", "Step"))]
        return lines[:4], tokens

    def check_connection(self):
        try:
            content, tokens = self._call_llm(
                "Eres un asistente útil.",
                "Responde solo con: OK",
                max_tokens=10
            )
            return True, f"{content} (modelo: {self._get_current_model()})"
        except Exception as e:
            return False, str(e)

    def get_status(self):
        return {
            "api_keys": len(self.api_keys),
            "current_key": self.current_key_index + 1,
            "models": self.models,
            "current_model": self._get_current_model(),
            "last_request": self.last_request_time
        }


def load_literouter_config(base_dir):
    config_path = os.path.join(base_dir, "Config", "plumloom_engine.yaml")
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    lit_cfg = config.get("endpoints", {}).get("literouter", {})
    api_keys_file = os.path.join(base_dir, "CONTEXTO", "apis literouter.txt")

    api_keys = []
    if os.path.exists(api_keys_file):
        with open(api_keys_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line.startswith("key") and "=" in line:
                    key = line.split("=", 1)[1].strip()
                    if key:
                        api_keys.append(key)

    if not api_keys:
        api_keys = [lit_cfg.get("api_key", "")]

    models = lit_cfg.get("models", [
        "deepseek-v3.2:free",
        "deepseek-v4-flash:free",
        "glm-5.3-flash:free",
        "qwen3.5:free"
    ])

    return {
        "api_keys": api_keys,
        "models": models,
        "timeout": lit_cfg.get("timeout", 30)
    }