import base64 
import json 
import os
import re 
from pathlib import Path 
from typing import Dict ,List ,Optional ,Union 

try :
    import ollama 
    OLLAMA_AVAILABLE =True 
except ImportError :
    OLLAMA_AVAILABLE =False 
    print ("Ollama not installed. Run: pip install ollama")

DEFAULT_HOST = os.environ.get("OLLAMA_HOST", "https://ollama.sirblob.co")
DEFAULT_MODEL = "ministral-3:latest"

DEFAULT_PROMPT ="""Analyze this image and identify ALL logos, brand names, and company names visible.

For each logo or brand you detect, provide:
1. The company/brand name
2. Confidence level (high, medium, low)
3. Location in image (top-left, center, bottom-right, etc.)
4. Product category if identifiable

Return your response as a valid JSON object with this exact structure:
{
    "logos_detected": [
        {
            "brand": "Company Name",
            "confidence": "high",
            "location": "center",
            "category": "technology"
        }
    ],
    "total_count": 1,
    "description": "Brief description of what's in the image"
}

If no logos are found, return:
{
    "logos_detected": [],
    "total_count": 0,
    "description": "Description of image with no visible logos"
}

IMPORTANT: Return ONLY the JSON object, no other text."""

class OllamaLogoDetector :
    def __init__ (self ,
    model :str =DEFAULT_MODEL ,
    host :str =DEFAULT_HOST ):
        if not OLLAMA_AVAILABLE :
            raise RuntimeError ("Ollama not installed. Run: pip install ollama")

        self .model =model 
        self .host =host 
        self .client =ollama .Client (host =host )

        try :
            models =self .client .list ()
            model_names =[m ['name']for m in models .get ('models',[])]

            model_base =model .split (':')[0 ]
            if not any (model_base in name for name in model_names ):
                print (f"Model '{model }' not found. Available models: {model_names }")
                print (f"Pulling {model }...")
                self .client .pull (model )
                print (f"Model {model } ready!")
            else :
                print (f"Using Ollama model: {model }")
        except Exception as e :
            print (f"Could not verify model: {e }")
            print ("Make sure Ollama is running: ollama serve")

    def detect_from_file (self ,
    image_path :str ,
    prompt :Optional [str ]=None )->Dict :
        path =Path (image_path )
        if not path .exists ():
            raise FileNotFoundError (f"Image not found: {image_path }")

        with open (path ,'rb')as f :
            image_data =base64 .b64encode (f .read ()).decode ('utf-8')

        return self ._analyze_image (image_data ,prompt )

    def detect_from_bytes (self ,
    image_bytes :bytes ,
    prompt :Optional [str ]=None )->Dict :
        image_data =base64 .b64encode (image_bytes ).decode ('utf-8')
        return self ._analyze_image (image_data ,prompt )

    def detect_from_numpy (self ,
    image_array ,
    prompt :Optional [str ]=None )->Dict :
        import cv2 

        success ,buffer =cv2 .imencode ('.jpg',image_array )
        if not success :
            raise ValueError ("Failed to encode image")

        return self .detect_from_bytes (buffer .tobytes (),prompt )

    def _analyze_image (self ,
    image_base64 :str ,
    prompt :Optional [str ]=None )->Dict :
        if prompt is None :
            prompt =DEFAULT_PROMPT 

        try :
            response =self .client .chat (
            model =self .model ,
            messages =[{
            'role':'user',
            'content':prompt ,
            'images':[image_base64 ]
            }],
            options ={
            'temperature':0.1 ,
            }
            )

            content =response ['message']['content']
            return self ._parse_response (content )

        except Exception as e :
            return {
            "logos_detected":[],
            "total_count":0 ,
            "error":str (e ),
            "raw_response":None 
            }

    def _parse_response (self ,content :str )->Dict :
        try :
            return json .loads (content )
        except json .JSONDecodeError :
            pass 

        json_patterns =[
        r'```json\s*([\s\S]*?)\s*```',
        r'```\s*([\s\S]*?)\s*```',
        r'\{[\s\S]*\}'
        ]

        for pattern in json_patterns :
            match =re .search (pattern ,content )
            if match :
                try :
                    json_str =match .group (1 )if '```'in pattern else match .group (0 )
                    return json .loads (json_str )
                except json .JSONDecodeError :
                    continue 

        return {
        "logos_detected":[],
        "total_count":0 ,
        "raw_response":content ,
        "parse_error":"Could not extract valid JSON from response"
        }

    def get_brands_list (self ,result :Dict )->List [str ]:
        logos =result .get ("logos_detected",[])
        return [logo .get ("brand","Unknown")for logo in logos ]
