import os
import json
import re
import mimetypes
import logging
import pypdf
from pathlib import Path

# Suppress noisy "Ignoring wrong pointing object" warnings from pypdf
logging.getLogger("pypdf").setLevel(logging.ERROR)
from bs4 import BeautifulSoup
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.messages import BinaryContent
from dotenv import load_dotenv
from VerifierApp.localisation_module import t, TRANSLATIONS
from VerifierApp.core.models import get_model_capabilities

def _t_en(key: str, **kwargs) -> str:
    """Always return the English translation for LLM-facing prompt strings."""
    text = TRANSLATIONS["en"].get(key, key)
    if kwargs:
        text = text.format(**kwargs)
    return text

load_dotenv()

def preprocess_pdf(file_path: Path) -> str:
    """Extract plain text from a PDF file using pypdf.

    Each page's text is joined with newlines.  Returns an empty string if the
    PDF contains no extractable text (e.g. scanned images without OCR).
    """
    try:
        reader = pypdf.PdfReader(str(file_path))
        pages = []
        for page in reader.pages:
            text = page.extract_text() or ""
            if text.strip():
                pages.append(text.strip())
        return "\n\n".join(pages)
    except Exception as e:
        print(f"PDF extraction failed for {file_path.name}: {e}")
        return ""


def preprocess_text(file_path: Path):
    content = file_path.read_text(encoding='utf-8')
    if file_path.suffix == '.html':
        soup = BeautifulSoup(content, 'html.parser')
        for tag in soup(["script", "style", "nav", "footer"]):
            tag.decompose()
        return soup.get_text(separator=' ', strip=True)
    return content

def extract_clean_json(raw_text: str) -> str:
    """Extract the first valid JSON array (or object) from a raw LLM response.

    Returns a JSON string.  Falls back to '[]' if nothing parseable is found
    so downstream code always receives valid JSON.
    """
    if not raw_text:
        return "[]"

    raw_text = raw_text.strip()

    start_idx = raw_text.find('[')
    if start_idx != -1:
        stack = 0
        in_string = False
        escape = False
        for i in range(start_idx, len(raw_text)):
            char = raw_text[i]
            if char == '"' and not escape:
                in_string = not in_string
            escape = (char == '\\' and not escape)
            if not in_string:
                if char == '[':
                    stack += 1
                elif char == ']':
                    stack -= 1
                    if stack == 0:
                        candidate = raw_text[start_idx: i + 1]
                        try:
                            json.loads(candidate)
                            return candidate
                        except json.JSONDecodeError:
                            continue

    match = re.search(r'\[\s*\{.*?\}\s*\]', raw_text, re.DOTALL)
    if match:
        candidate = match.group(0)
        try:
            json.loads(candidate)
            return candidate
        except json.JSONDecodeError:
            pass

    obj_start = raw_text.find('{')
    if obj_start != -1:
        stack = 0
        in_string = False
        escape = False
        for i in range(obj_start, len(raw_text)):
            char = raw_text[i]
            if char == '"' and not escape:
                in_string = not in_string
            escape = (char == '\\' and not escape)
            if not in_string:
                if char == '{':
                    stack += 1
                elif char == '}':
                    stack -= 1
                    if stack == 0:
                        candidate = raw_text[obj_start: i + 1]
                        try:
                            obj = json.loads(candidate)
                            return json.dumps([obj])
                        except json.JSONDecodeError:
                            continue

    print("Extractor response wasn't a valid JSON: " + raw_text)
    return "[]"

async def run_extraction(input_path_str: str, output_path_str: str, model_path_str: str = None,
                         custom_prompt: str = None, llm_model_name: str = "qwen3.5-122b"):
    input_dir = Path(input_path_str)
    output_path = Path(output_path_str)

    try:
        gemini_token = os.getenv('GEMINI_API_TOKEN')
        is_gemini = llm_model_name.startswith("gemini") and gemini_token

        llm_model = OpenAIChatModel(
            llm_model_name,
            provider=OpenAIProvider(
                base_url='https://generativelanguage.googleapis.com/v1beta/openai/' if is_gemini
                         else 'https://llm.ai.e-infra.cz/v1',
                api_key=gemini_token if is_gemini else os.getenv('E_INFRA_API_TOKEN'),
            ),
        )
        agent = Agent(llm_model, system_prompt=custom_prompt)
    except Exception as e:
        print(t("extractor_init_error", model=llm_model_name, error=e))
        return None

    if output_path.is_dir() or (not output_path.suffix):
        output_file = output_path / "phenotypes.json"
        output_file.parent.mkdir(parents=True, exist_ok=True)
    else:
        output_file = output_path
        output_file.parent.mkdir(parents=True, exist_ok=True)

    if not input_dir.exists():
        print(t("extractor_dir_missing", path=input_dir))
        return None

    prompt_parts = [_t_en("extractor_prompt_intro")]

    if model_path_str:
        model_path = Path(model_path_str)
        if model_path.exists():
            try:
                model_content = model_path.read_text(encoding='utf-8')
                prompt_parts.append(
                    _t_en("extractor_model_context_header", name=model_path.name)
                    + model_content + "\n"
                    + _t_en("extractor_model_context_footer")
                )
            except Exception as e:
                print(t("extractor_model_load_error", error=e))

    caps = get_model_capabilities(llm_model_name)

    for file_path in input_dir.iterdir():
        if file_path.suffix.lower() in ['.html', '.txt']:
            text = preprocess_text(file_path)
            prompt_parts.append(_t_en("extractor_text_section", name=file_path.name, text=text))
        elif file_path.suffix.lower() == '.pdf':
            if caps["pdf"]:
                print(f"Sending {file_path.name} as native PDF to {llm_model_name}")
                prompt_parts.append(BinaryContent(data=file_path.read_bytes(), media_type='application/pdf'))
            else:
                print(f"Extracting text from {file_path.name} ({llm_model_name} does not support native PDF)")
                pdf_text = preprocess_pdf(file_path)
                if pdf_text.strip():
                    prompt_parts.append(_t_en("extractor_text_section", name=file_path.name, text=pdf_text))
                else:
                    print(f"Skipping {file_path.name}: no extractable text found in PDF.")
        elif file_path.suffix.lower() in ['.png', '.jpg', '.jpeg']:
            if caps["vision"]:
                mime_type, _ = mimetypes.guess_type(str(file_path))
                prompt_parts.append(BinaryContent(data=file_path.read_bytes(), media_type=mime_type or 'image/jpeg'))
            else:
                print(f"Skipping {file_path.name}: {llm_model_name} does not support image input.")

    try:
        result = await agent.run(prompt_parts)

        if hasattr(result, 'data'):
            raw_output = str(result.data).strip()
        elif hasattr(result, 'all_messages'):
            raw_output = str(result.all_messages()[-1].parts[-1].content).strip()
        else:
            raw_output = str(result).strip()

        clean_json = extract_clean_json(raw_output)
        output_file.write_text(clean_json, encoding='utf-8')

        usage = {}
        try:
            u = result.usage()
            usage = {
                "request_tokens":  u.input_tokens,
                "response_tokens": u.output_tokens,
                "total_tokens":    u.total_tokens,
            }
        except Exception:
            pass

        return {"phenotypes": clean_json, "usage": usage}

    except Exception as e:
        print(t("extractor_api_error", model=llm_model_name, error=e))
        return None