import gradio as gr
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
import numpy as np
import logging
from dataclasses import dataclass
from typing import Dict, List, Tuple, Optional
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import json

# ----------------------------------
# Logging
# ----------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ----------------------------------
# Config
# ----------------------------------
@dataclass
class AppConfig:
    model_name: str = "dejanseo/link-prediction"
    max_length: int = 512
    doc_stride: int = 128
    device: str = "cuda" if torch.cuda.is_available() else "cpu"

# ----------------------------------
# Load model/tokenizer
# ----------------------------------
config = AppConfig()
logger.info(f"Loading model: {config.model_name} on {config.device}")

model = AutoModelForTokenClassification.from_pretrained(config.model_name)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)
model.to(config.device)
model.eval()

logger.info("Model loaded successfully.")

# ----------------------------------
# Inference helpers
# ----------------------------------
def windowize_inference(
    plain_text: str, tokenizer: AutoTokenizer, max_length: int, doc_stride: int
) -> List[Dict]:
    """Slice long text into overlapping windows for inference."""
    specials = tokenizer.num_special_tokens_to_add(pair=False)
    cap = max_length - specials
    full_encoding = tokenizer(
        plain_text, add_special_tokens=False, return_offsets_mapping=True, truncation=False
    )
    temp_tokenization = tokenizer(plain_text, truncation=False)
    full_word_ids = temp_tokenization.word_ids(batch_index=0)

    windows_data = []
    step = max(cap - doc_stride, 1)
    start_token_idx = 0
    total_tokens = len(full_encoding["input_ids"])

    if total_tokens == 0 and len(plain_text) > 0:
        logger.warning("Tokenizer produced 0 tokens for a non-empty string.")
        return []

    while start_token_idx < total_tokens:
        end_token_idx = min(start_token_idx + cap, total_tokens)
        ids_slice = full_encoding["input_ids"][start_token_idx:end_token_idx]
        offsets_slice = full_encoding["offset_mapping"][start_token_idx:end_token_idx]

        word_ids_slice = []
        current_token = 0
        for i, wid in enumerate(full_word_ids):
            if temp_tokenization.token_to_chars(i) is not None:
                if current_token >= start_token_idx and current_token < end_token_idx:
                    word_ids_slice.append(wid)
                current_token += 1

        input_ids = tokenizer.build_inputs_with_special_tokens(ids_slice)
        attention_mask = [1] * len(input_ids)
        padding_length = max_length - len(input_ids)
        input_ids.extend([tokenizer.pad_token_id] * padding_length)
        attention_mask.extend([0] * padding_length)

        window_offset_mapping = [(0, 0)] + offsets_slice + [(0, 0)]
        window_offset_mapping += [(0, 0)] * padding_length

        window_word_ids = [None] + word_ids_slice + [None]
        window_word_ids += [None] * padding_length

        windows_data.append({
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "word_ids": window_word_ids[:max_length],
            "offset_mapping": window_offset_mapping[:max_length],
        })
        if end_token_idx >= total_tokens:
            break
        start_token_idx += step
    return windows_data


def classify_text(text: str, threshold_percent: float) -> Tuple[str, Optional[str]]:
    """Classify link tokens with windowing. Returns (html, warning)."""
    if not text.strip():
        return "", "Input text is empty."

    windows = windowize_inference(text, tokenizer, config.max_length, config.doc_stride)
    if not windows:
        return "", "Could not generate any windows for processing."

    char_link_probabilities = np.zeros(len(text), dtype=np.float32)

    with torch.no_grad():
        for window in windows:
            inputs = {
                'input_ids': window['input_ids'].unsqueeze(0).to(config.device),
                'attention_mask': window['attention_mask'].unsqueeze(0).to(config.device)
            }
            outputs = model(**inputs)
            probabilities = torch.softmax(outputs.logits, dim=-1).squeeze(0)
            link_probs = probabilities[:, 1].cpu().numpy()

            for i, offset in enumerate(window['offset_mapping']):
                if isinstance(offset, (list, tuple)) and len(offset) == 2:
                    start, end = offset
                    if window['word_ids'][i] is not None and start < end:
                        char_link_probabilities[start:end] = np.maximum(
                            char_link_probabilities[start:end], link_probs[i]
                        )

    final_threshold = threshold_percent / 100.0

    full_encoding = tokenizer(text, return_offsets_mapping=True, truncation=False)
    word_ids = full_encoding.word_ids(batch_index=0)
    offsets = full_encoding['offset_mapping']

    word_max_prob_map: Dict[int, float] = {}
    word_char_spans: Dict[int, List[int]] = {}

    for i, word_id in enumerate(word_ids):
        if word_id is not None and i < len(offsets):
            start_char, end_char = offsets[i]
            if start_char < end_char:
                current_token_max_prob = np.max(char_link_probabilities[start_char:end_char]) if start_char < len(char_link_probabilities) else 0.0

                if word_id not in word_max_prob_map:
                    word_max_prob_map[word_id] = current_token_max_prob
                    word_char_spans[word_id] = [start_char, end_char]
                else:
                    word_max_prob_map[word_id] = max(word_max_prob_map[word_id], current_token_max_prob)
                    word_char_spans[word_id][1] = end_char

    highlight_candidates: Dict[int, float] = {}
    for word_id, max_prob in word_max_prob_map.items():
        if max_prob >= final_threshold:
            highlight_candidates[word_id] = max_prob

    max_highlight_prob = max(highlight_candidates.values()) if highlight_candidates else 0.0

    html_parts, current_char = [], 0
    sorted_word_ids = sorted(word_char_spans.keys(), key=lambda k: word_char_spans[k][0])

    for word_id in sorted_word_ids:
        start_char, end_char = word_char_spans[word_id]

        if start_char > current_char:
            html_parts.append(text[current_char:start_char])

        word_text = text[start_char:end_char]

        if word_id in highlight_candidates:
            word_prob = highlight_candidates[word_id]
            normalized_opacity = (word_prob / max_highlight_prob) * 0.9 + 0.1 if max_highlight_prob > 0 else 1.0

            html_parts.append(
                f"<span style='background-color: #D4EDDA; color: #155724; "
                f"padding: 0.1em 0.2em; border-radius: 0.2em; opacity: {normalized_opacity:.2f};' "
                f"title='Link Probability: {word_prob:.1%}'>{word_text}</span>"
            )
        else:
            html_parts.append(word_text)
        current_char = end_char

    if current_char < len(text):
        html_parts.append(text[current_char:])

    return "".join(html_parts), None


# ----------------------------------
# Gradio Interface
# ----------------------------------
def predict(text: str, threshold: float) -> str:
    """Main prediction function for Gradio."""
    html, warning = classify_text(text, threshold)
    if warning:
        return f"<p style='color: orange;'>{warning}</p>"
    return html


def classify_text_json(text: str, threshold_percent: float) -> Dict:
    """
    Returns structured JSON with word-level predictions.
    Frontend can render this however it wants.
    """
    if not text.strip():
        return {"error": "Input text is empty.", "words": []}

    windows = windowize_inference(text, tokenizer, config.max_length, config.doc_stride)
    if not windows:
        return {"error": "Could not process text.", "words": []}

    char_link_probabilities = np.zeros(len(text), dtype=np.float32)

    with torch.no_grad():
        for window in windows:
            inputs = {
                'input_ids': window['input_ids'].unsqueeze(0).to(config.device),
                'attention_mask': window['attention_mask'].unsqueeze(0).to(config.device)
            }
            outputs = model(**inputs)
            probabilities = torch.softmax(outputs.logits, dim=-1).squeeze(0)
            link_probs = probabilities[:, 1].cpu().numpy()

            for i, offset in enumerate(window['offset_mapping']):
                if isinstance(offset, (list, tuple)) and len(offset) == 2:
                    start, end = offset
                    if window['word_ids'][i] is not None and start < end:
                        char_link_probabilities[start:end] = np.maximum(
                            char_link_probabilities[start:end], link_probs[i]
                        )

    final_threshold = threshold_percent / 100.0

    full_encoding = tokenizer(text, return_offsets_mapping=True, truncation=False)
    word_ids = full_encoding.word_ids(batch_index=0)
    offsets = full_encoding['offset_mapping']

    word_max_prob_map: Dict[int, float] = {}
    word_char_spans: Dict[int, List[int]] = {}

    for i, word_id in enumerate(word_ids):
        if word_id is not None and i < len(offsets):
            start_char, end_char = offsets[i]
            if start_char < end_char:
                current_token_max_prob = np.max(char_link_probabilities[start_char:end_char]) if start_char < len(char_link_probabilities) else 0.0

                if word_id not in word_max_prob_map:
                    word_max_prob_map[word_id] = current_token_max_prob
                    word_char_spans[word_id] = [start_char, end_char]
                else:
                    word_max_prob_map[word_id] = max(word_max_prob_map[word_id], current_token_max_prob)
                    word_char_spans[word_id][1] = end_char

    # Build structured output
    words_output = []
    sorted_word_ids = sorted(word_char_spans.keys(), key=lambda k: word_char_spans[k][0])

    for word_id in sorted_word_ids:
        start_char, end_char = word_char_spans[word_id]
        word_text = text[start_char:end_char]
        prob = float(word_max_prob_map[word_id])
        is_link = prob >= final_threshold

        words_output.append({
            "text": word_text,
            "start": start_char,
            "end": end_char,
            "probability": round(prob, 4),
            "is_link_candidate": is_link
        })

    return {
        "error": None,
        "threshold": threshold_percent,
        "words": words_output,
        "original_text": text
    }


def predict_json(text: str, threshold: float) -> str:
    """Returns JSON string with structured word-level predictions."""
    import json
    result = classify_text_json(text, threshold)
    return json.dumps(result)


# Build the interface
with gr.Blocks(title="LinkBERT by DEJAN AI") as demo:
    gr.Markdown("# LinkBERT")
    gr.Markdown("Predict natural link placement in plain text.")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Input Text",
                placeholder="Paste your text here...",
                lines=8,
                value="DEJAN AI is the world's leading AI SEO agency. This tool showcases the capability of our latest link prediction model called LinkBERT."
            )
            threshold_slider = gr.Slider(
                minimum=0,
                maximum=100,
                value=70,
                step=1,
                label="Link Probability Threshold (%)"
            )
            submit_btn = gr.Button("Classify Text", variant="primary")
        
        with gr.Column():
            output_html = gr.HTML(label="Results")
    
    submit_btn.click(
        fn=predict,
        inputs=[text_input, threshold_slider],
        outputs=output_html
    )


# === Create FastAPI app with custom routes, then mount Gradio ===
class ClassifyRequest(BaseModel):
    text: str
    threshold: float = 70.0


# Create FastAPI app FIRST
app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Add custom API routes BEFORE mounting Gradio
@app.post("/api/classify")
async def api_classify(request: ClassifyRequest):
    """Direct REST endpoint"""
    result = classify_text_json(request.text, request.threshold)
    return result


@app.get("/api/health")
async def health():
    return {"status": "ok", "model": config.model_name}


@app.get("/")
async def root():
    """Redirect root to Gradio UI"""
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/gradio")


# Mount Gradio AFTER defining routes
app = gr.mount_gradio_app(app, demo, path="/gradio")


# Launch with uvicorn
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)