from __future__ import annotations

import datetime as dt
import json
import uuid
from dataclasses import dataclass
from typing import Dict, List, Optional

import giskard
from giskard.core.suite import Suite
from giskard.llm import set_llm_model
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_google_vertexai import ChatVertexAI

from app.core.config import get_settings
from app.models.firewall import SecurityCheckRequest
from app.models.scan import FirewallGroupResult, IssueDetail, ScanRequest, ScanResult
from app.services import credits
from app.services.firewall import run_firewall_check, FirewallError
from app.services.llm_client import LLMContentFilteredError, invoke_chat


@dataclass
class ScanContext:
    user_id: str
    scan_id: str
    categories: List[str]
    title: str
    description: str


class ScannerError(Exception):
    ...


def generate_title_and_description(prompt: str) -> Dict[str, str]:
    settings = get_settings()
    messages = [
        {
            "role": "system",
            "content": "You label prompts with a concise title and 1-2 sentence description",
        },
        {
            "role": "user",
            "content": f"Prompt:\n{prompt}\n\nReturn JSON with 'title' and 'description' fields.",
        },
    ]
    try:
        response = invoke_chat(messages, temperature=0.3)
        return {
            "title": response.get("title", "LLM Prompt Analysis"),
            "description": response.get(
                "description", "Analysis of LLM behavior and responses for the given prompt"
            ),
        }
    except LLMContentFilteredError:
        return {
            "title": "LLM Prompt Analysis",
            "description": "Content could not be analyzed due to safety restrictions.",
        }


def build_llm_chain(prompt: str, context_str: Optional[str]) -> LLMChain:
    set_llm_model("gemini-2.0-flash-lite")
    # If context is provided, include it via a partial variable so only "question" is required at runtime
    has_context = bool(context_str)
    context_block = "Context:\n{context}\n\n" if has_context else ""
    template = f"{prompt}\n\n{context_block}Question:\n{{question}}\n\nYour answer:"

    settings = get_settings()
    llm = ChatVertexAI(
        model_name=settings.vertex_ai.model,
        temperature=0,
        max_output_tokens=None,
        location=settings.vertex_ai.location,
        project=settings.vertex_ai.project_id,
    )
    if has_context:
        prompt_template = PromptTemplate(
            template=template,
            input_variables=["question"],
            partial_variables={"context": context_str or ""},
        )
    else:
        prompt_template = PromptTemplate(template=template, input_variables=["question"])
    return LLMChain(llm=llm, prompt=prompt_template, output_key="result")


def run_giskard_scan(chain: LLMChain, categories: List[str], context: ScanContext):
    model = giskard.Model(
        model=chain,
        model_type="text_generation",
        name=context.title,
        description=context.description,
        feature_names=["question"],
    )

    report = giskard.scan(model, only=categories)

    # Optionally persist and execute the generated test suite (disabled by default in Cloud Run)
    settings = get_settings()
    path_to_tests = None
    try:
        if getattr(settings, "scanner_save_tests", False):
            test_suite_name = "_".join(categories)
            test_suite = report.generate_test_suite(test_suite_name)
            path_to_tests = f"{context.user_id}/{context.scan_id}/{test_suite_name}"
            test_suite.save(folder=path_to_tests)
            Suite.load(folder=path_to_tests).run()
    except Exception:
        # Do not fail the scan if suite persistence/execution is unavailable
        path_to_tests = None

    return report, path_to_tests


def _format_report(report) -> Dict:
    formatted = {
        "report_attributes": giskard.scan_utils.to_dict(report),
        "issues": [],
    }
    for issue in report.issues:
        formatted["issues"].append({
            "id": issue.id,
            "group": getattr(issue, "group", ""),
            "level": getattr(issue, "level", ""),
            "scan_examples": getattr(issue, "scan_examples", {}),
        })
    return formatted


def _extract_conversations(issue: Dict) -> List[Dict[str, str]]:
    examples = issue.get("scan_examples", {}).get("examples", [])
    conversations = []
    for example in examples:
        conversation = example.get("Conversation", "")
        if not conversation:
            continue
        parts = conversation.split("\n\nAGENT: ")
        if len(parts) < 2:
            continue
        user_part = parts[0].replace("USER: ", "").strip()
        agent_part = parts[1].strip()
        if user_part and agent_part:
            conversations.append({"user": user_part, "assistant": agent_part})
    return conversations


def execute_scan(request: ScanRequest, user_id: str) -> ScanResult:
    if not request.prompt:
        raise ScannerError("Prompt is required")
    if not request.categories:
        raise ScannerError("At least one category is required")

    scan_id = uuid.uuid4().hex[:12]
    context = ScanContext(
        user_id=user_id,
        scan_id=scan_id,
        categories=request.categories,
        title=request.title or "",
        description=request.description or "",
    )

    if not context.title or not context.description:
        generated = generate_title_and_description(request.prompt)
        context.title = context.title or generated["title"]
        context.description = context.description or generated["description"]

    try:
        credits.deduct_scanner_credits(user_id, len(request.categories))
    except credits.InsufficientCreditsError as exc:
        raise ScannerError("Insufficient credits for scanner run") from exc

    # "model" field on request is not a context string; keep template stable and only vary "question".
    chain = build_llm_chain(request.prompt, None)
    report, _ = run_giskard_scan(chain, request.categories, context)
    formatted_report = _format_report(report)

    issue_details: List[IssueDetail] = []
    for issue in formatted_report.get("issues", []):
        issue_details.append(
            IssueDetail(
                id=issue.get("id", ""),
                group=issue.get("group", ""),
                level=issue.get("level", ""),
                description=None,
            )
        )

        conversations = _extract_conversations(issue)
        for convo in conversations:
            try:
                run_firewall_check(
                    SecurityCheckRequest(
                        system_prompt=request.prompt,
                        user_input=convo["user"],
                        response=convo["assistant"],
                        security_checks=request.categories,
                    ),
                    user_id=user_id,
                    source="scanner",
                    event_id=f"scan-{scan_id}",
                )
            except FirewallError:
                continue

    return ScanResult(
        issues=issue_details,
        firewall_results=[],
        timestamp=dt.datetime.now(dt.timezone.utc),
        credits_used=len(request.categories) * credits.CREDIT_COSTS["scanner_category"],
    )
