import ast
import csv
import json
import logging
import time
from pathlib import Path
from typing import Dict, Tuple

import gradio as gr
import numpy as np
import requests
from huggingface_hub import hf_hub_download
from PIL import Image

# Configure Logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(message)s")
logger = logging.getLogger("PixAITagger")

logging.getLogger().setLevel(logging.WARNING)
logging.getLogger("PixAITagger").setLevel(logging.INFO)

# Constants
MODEL_REPO = "deepghs/pixai-tagger-v0.9-onnx"
MODEL_FILENAME = "model.onnx"
TAGS_FILENAME = "selected_tags.csv"
INPUT_SIZE = 448

# Helper for URL fetch (included directly to ensure standalone functionality)
def fetch_image_from_url(url: str) -> Image.Image:
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        from io import BytesIO
        return Image.open(BytesIO(response.content))
    except Exception as e:
        logger.error(f"Failed to fetch image: {e}")
        return None

class HybridEngine:
    """
    Handles inference with dynamic backend selection.
    Priority: OpenVINO INT8 -> OpenVINO FP32 -> ONNX Runtime
    """

    def __init__(self, model_path: str, use_int8: bool = True):
        self.model_path = model_path
        self.session = None
        self.use_openvino = False
        self.provider_name = None
        self.use_int8 = use_int8

        self._init_backend()

    def _init_backend(self):
        # --- Attempt 1 & 2: OpenVINO (INT8 or FP32) ---
        try:
            import openvino as ov
            logger.info("Engine: OpenVINO available, reading model...")
            core = ov.Core()
            model = core.read_model(self.model_path)
            
            # Logic for INT8 Quantization
            if self.use_int8:
                try:
                    import nncf
                    logger.info("Engine: Compressing weights to INT8 using NNCF...")
                    model = nncf.compress_weights(model)
                    self.provider_name = "OpenVINO (INT8 Weights)"
                except ImportError:
                    logger.warning("Engine: NNCF not installed. Falling back to FP32.")
                    self.provider_name = "OpenVINO (FP32 - NNCF missing)"
                except Exception as e:
                    logger.warning(f"Engine: NNCF Compression failed ({e}). Falling back to FP32.")
                    self.provider_name = "OpenVINO (FP32 - Compression error)"
            else:
                self.provider_name = "OpenVINO (FP32)"

            # Compile
            self.session = core.compile_model(model, "CPU")
            self.use_openvino = True
            logger.info(f"Engine: Success using {self.provider_name}")
            return

        except Exception as e:
            logger.warning(f"Engine: OpenVINO initialization failed ({e}). Falling back to ONNX Runtime.")

        # --- Attempt 3: ONNX Runtime (Fallback) ---
        try:
            import onnxruntime as ort
            sess_options = ort.SessionOptions()
            sess_options.log_severity_level = 3

            self.session = ort.InferenceSession(
                self.model_path,
                sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )
            self.use_openvino = False
            self.provider_name = f"ONNX Runtime ({self.session.get_providers()[0]})"
            logger.info("Engine: Using ONNX Runtime backend")
        except Exception as e:
            logger.error(f"Engine: FATAL - ONNX Runtime also failed: {e}")
            self.provider_name = "Error: No backend available"
            self.session = None

    def run(self, input_data: np.ndarray, expected_dim: int):
        if not self.session:
            raise RuntimeError("Engine not initialized")

        if self.use_openvino:
            results = self.session(input_data)
            outputs = list(results.values())
        else:
            input_name = self.session.get_inputs()[0].name
            outputs = self.session.run(None, {input_name: input_data})

        # Pick output with expected class dimension
        for out in outputs:
            if out.shape[1] == expected_dim:
                return out[0], self.provider_name

        # Fallback: largest output
        out = max(outputs, key=lambda x: x.shape[1])
        return out[0], self.provider_name


class PixAITagger:
    def __init__(self):
        self.model_path = None
        self.tags_list = []
        self._load_resources()
        
        # State tracking for engine reloading
        self.engine = None
        self.current_int8_mode = None 

    def _load_resources(self):
        logger.info(f"Downloading resources from {MODEL_REPO}...")
        self.model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)

        try:
            tags_path = hf_hub_download(repo_id=MODEL_REPO, filename=TAGS_FILENAME)
            self._load_tags_csv(Path(tags_path))
        except Exception as e:
            raise FileNotFoundError(f"Could not load tags file: {e}")

    def _load_tags_csv(self, csv_path: Path):
        self.tags_list = []
        with csv_path.open("r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                try:
                    idx = int(row.get("id"))
                    name = row.get("name")
                    category = int(row.get("category", 0))
                    ips_raw = row.get("ips", "[]")
                    ips = []
                    if ips_raw and ips_raw != "[]":
                        try:
                            ips = json.loads(ips_raw)
                        except Exception:
                            try:
                                ips = ast.literal_eval(ips_raw)
                            except Exception:
                                pass
                    self.tags_list.append(
                        {"id": idx, "name": name, "is_char": category == 4, "ips": [str(ip) for ip in ips]}
                    )
                except ValueError:
                    continue

        self.tags_list.sort(key=lambda x: x["id"])

        self.id_to_tag = {}
        self.char_indices = []
        self.gen_indices = []
        self.mapping = {}

        for item in self.tags_list:
            idx = item["id"]
            name = item["name"]
            self.id_to_tag[idx] = name

            if item["is_char"]:
                self.char_indices.append(idx)
                if item["ips"]:
                    self.mapping[name] = item["ips"]
            else:
                self.gen_indices.append(idx)

        self.num_classes = len(self.tags_list)
        logger.info(f"Loaded {self.num_classes} tags.")

    def preprocess(self, image: Image.Image) -> np.ndarray:
        if image.mode != "RGB":
            image = image.convert("RGB")
        image = image.resize((INPUT_SIZE, INPUT_SIZE), Image.BICUBIC)
        img = np.array(image).astype(np.float32) / 255.0
        img = (img - 0.5) / 0.5
        img = img.transpose(2, 0, 1)
        return np.expand_dims(img, 0)

    def predict(
        self, image: Image.Image, gen_threshold: float, char_threshold: float, resolve_mapping: bool, use_int8_weights: bool
    ) -> Tuple[Dict, Dict, str, str, float]:
        
        # Reload engine if the INT8 preference changed or engine doesn't exist
        if self.engine is None or self.current_int8_mode != use_int8_weights:
            logger.info(f"Reloading engine. New mode INT8: {use_int8_weights}")
            self.engine = HybridEngine(str(self.model_path), use_int8=use_int8_weights)
            self.current_int8_mode = use_int8_weights

        input_tensor = self.preprocess(image)
    
        infer_start = time.time()
        logits, provider_name = self.engine.run(input_tensor, self.num_classes)
        infer_time = time.time() - infer_start
    
        # Sigmoid
        probs = 1 / (1 + np.exp(-logits))
    
        # General Tags
        gen_tags = {}
        for idx in self.gen_indices:
            if idx < len(probs):
                score = float(probs[idx])
                if score >= gen_threshold:
                    gen_tags[self.id_to_tag[idx]] = score
    
        # Character Tags & IPs
        char_tags = {}
        detected_ips = set()
    
        for idx in self.char_indices:
            if idx < len(probs):
                score = float(probs[idx])
                if score >= char_threshold:
                    name = self.id_to_tag[idx]
                    char_tags[name] = score
    
                    if resolve_mapping and name in self.mapping:
                        for ip in self.mapping[name]:
                            detected_ips.add(ip)
    
        gen_tags = dict(sorted(gen_tags.items(), key=lambda x: x[1], reverse=True))
        char_tags = dict(sorted(char_tags.items(), key=lambda x: x[1], reverse=True))
    
        ip_text = ", ".join(sorted(list(detected_ips))) if detected_ips else ""
    
        return gen_tags, char_tags, ip_text, provider_name, infer_time


# --- UI Setup ---

tagger_instance = None

def get_tagger():
    global tagger_instance
    if tagger_instance is None:
        tagger_instance = PixAITagger()
    return tagger_instance

def init_app():
    """Warms up the model loader."""
    get_tagger()
    return None

def run_inference(image, gen_thresh, char_thresh, resolve_mapping, use_int8):
    if image is None:
        return "", "", "", {}, {}, ""
    try:
        model = get_tagger()
        start_time = time.time()

        gen_tags, char_tags, ip_str, provider, infer_time = model.predict(
            image, gen_thresh, char_thresh, resolve_mapping, use_int8
        )

        char_str = ", ".join(char_tags.keys()).replace("_", " ")
        gen_str = ", ".join(gen_tags.keys()).replace("_", " ")

        if not resolve_mapping:
            ip_disp = ""
        elif not ip_str:
            ip_disp = ""
        else:
            ip_disp = ip_str.replace("_", " ")

        time_disp = f"- **Provider:** {provider} | **Inference time:** {infer_time:.4f}s"
        return char_str, ip_disp, gen_str, char_tags, gen_tags, time_disp
    except Exception as e:
        logger.error(f"Inference Error: {e}")
        return "", f"Error: {str(e)}", "", {}, {}, f"Error: {str(e)}"


with gr.Blocks(title="PixAI Tagger v0.9 ONNX") as demo:

    gr.Markdown(
        '<a href="https://huggingface.co/pixai-labs/pixai-tagger-v0.9" '
        'target="_blank" rel="noopener noreferrer">PixAI Tagger</a>'
        ' is an iteration on top of '
        '<a href="https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3" '
        'target="_blank" rel="noopener noreferrer">SmilingWolf/wd-eva02-large-tagger-v3</a>'
        ' with an updated dataset (2025-01).  \n'
        'It should be noted that PixAI Tagger may be worse in accuracy over eva02 large. See the PixAI page for details'
    )

    # Header Row
    with gr.Row(elem_classes=["container"]):
        with gr.Column(scale=1, elem_classes=["header-col"]):
            gr.Markdown("### Input")
        with gr.Column(scale=1, elem_classes=["header-col"]):
            gr.Markdown("### Configuration & Results")

    with gr.Row(elem_classes=["container"]):
        # LEFT COLUMN
        with gr.Column(scale=1):
            url_input = gr.Textbox(
                label="Enter Image URL (not all may work) or upload an image below",
                placeholder="https://example.com/image.jpg",
            )
            input_img = gr.Image(type="pil", label="", show_label=False, elem_classes=["image-container"])

        # RIGHT COLUMN - Controls and Outputs
        with gr.Column(scale=1):
            # Action Buttons
            with gr.Row():
                run_btn = gr.Button("Run (Image Upload)", variant="primary")
                url_btn = gr.Button("Run (URL)", variant="secondary")

            # Output Textboxes
            with gr.Group():
                with gr.Row(elem_id="d-row-container"):
                    char_box = gr.Textbox(label="Character Tags", interactive=False, buttons=["copy"])
                    ip_box = gr.Textbox(label="Character - Copyright Mapping", interactive=False, buttons=["copy"])

                gen_box = gr.Textbox(label="General Tags", interactive=False, buttons=["copy"])
                time_info = gr.Markdown(elem_id="time-display")

            # Configuration Section
            with gr.Group():
                with gr.Row(elem_id="d-row-container"):
                    char_slider = gr.Slider(0.0, 1.0, value=0.75, step=0.05, label="Character Threshold")
                    gen_slider = gr.Slider(0.0, 1.0, value=0.30, step=0.05, label="General Threshold")
                
                # Checkboxes Row
                with gr.Row():
                    map_checkbox = gr.Checkbox(value=True, label="Resolve Copyright Mapping")
                    int8_checkbox = gr.Checkbox(value=True, label="INT8 Weights")

            # Confidence Plots
            gr.Markdown("### Confidence Scores")
            char_plot = gr.Label(label="Character Probabilities", num_top_classes=50)
            gen_plot = gr.Label(label="General Probabilities", num_top_classes=500)

    gr.Markdown(
        'Model sourced from '
        '<a href="https://huggingface.co/deepghs/pixai-tagger-v0.9-onnx" '
        'target="_blank" rel="noopener noreferrer">deepghs/pixai-tagger-v0.9-onnx</a>.  \n'
        'OpenVINO™ will be used to accelerate CPU inference with ONNX CPUExecutionProvider as fallback.  \n'
        'INT8 weights option may improve inference times with a deviation of roughly +-0.005 in scores.' 
    )

    # Click Logic
    # Added int8_checkbox to inputs
    url_btn.click(fn=fetch_image_from_url, inputs=[url_input], outputs=[input_img]).then(
        fn=run_inference,
        inputs=[input_img, gen_slider, char_slider, map_checkbox, int8_checkbox],
        outputs=[char_box, ip_box, gen_box, char_plot, gen_plot, time_info],
    )

    run_btn.click(
        fn=run_inference,
        inputs=[input_img, gen_slider, char_slider, map_checkbox, int8_checkbox],
        outputs=[char_box, ip_box, gen_box, char_plot, gen_plot, time_info],
    )

    # Load logic
    demo.load(fn=init_app)

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Base(), 
        css="""
    * { box-sizing: border-box; }
    @media (max-width: 1022px) {
        #d-row-container {
            flex-direction: column !important;
        }
        #d-row-container > * {
            width: 100% !important;
        }
        #d-row-container .block {
            width: 100% !important;
        }
    }""",
    )