import logging
import os
from google import genai
from google.genai import types
from pydantic import BaseModel

from .distance_ocr import DistanceOCR

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

GEMINI_MODELS = ["gemini-2.5-flash", "gemini-2.5-flash-lite"]

DEFAULT_SYSTEM_INSTRUCTION = """
Extract the distance value overlaid on an industrial pipe inspection image.
Detect metric or imperial distance text. Treat comma as decimal separator.
Convert all values to standard decimal format.
Only convert number formatting, not units.
If the decimal part is unclear, round to nearest integer.
Output distance in decimal format without conversion.
If no distance is visible, output NONE.
No units. No extra text.
Examples:
1.25
2.0
NONE
"""


class GeminiImageDistanceConfig(BaseModel):
    model: str = "gemini-2.5-flash-lite"
    system_instruction: str = DEFAULT_SYSTEM_INSTRUCTION
    prompt: str = "extract distance"


class GeminiImageDistanceExtractor(DistanceOCR):
    def __init__(self, api_key: str | None = None, config: GeminiImageDistanceConfig | None = None):

        if api_key is None:
            self.api_key = os.getenv("GEMINI_API_KEY")
            if self.api_key is None:
                raise ValueError("No API key set. Pass in an api key or set os var GEMINI_API_KEY")
        else:
            self.api_key = api_key
    
        self.client = genai.Client(api_key=self.api_key)
        self.config = config or GeminiImageDistanceConfig()

    def get_image_distance(self, image_path: str, **kwargs) -> float | None:
        with open(image_path, "rb") as f:
            image_bytes = f.read()

        response = self.client.models.generate_content(
            model=self.config.model,
            config=types.GenerateContentConfig(
                system_instruction=self.config.system_instruction
            ),
            contents=[
                types.Part.from_bytes(
                    data=image_bytes,
                    mime_type="image/jpeg",
                ),
                self.config.prompt,
            ],
        )

        if response.text is None:
            raise ValueError(f"Received no response text from gemini. Response object: {response}")

        if response.text.upper() == "NONE":
            return None

        return float(response.text.strip())


def main() -> None:
    import argparse
    import os
    from dotenv import load_dotenv

    load_dotenv()

    parser = argparse.ArgumentParser(description="Extract distance from pipe inspection image using Gemini.")
    parser.add_argument("image", help="Path to the image file")
    args = parser.parse_args()

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("No api key set")
    extractor = GeminiImageDistanceExtractor(api_key=api_key)
    result = extractor.get_image_distance(args.image)
    print(result if result is not None else "NONE")


if __name__ == "__main__":
    main()

