import os
import cv2
import numpy as np
import duckdb
import glob
from shapely.geometry import Polygon
from sklearn.linear_model import LogisticRegression
import pickle
from datetime import datetime
from PIL import Image
import re
from PIL.ExifTags import TAGS
import hashlib

class SnowDetectionAgent:
    """
    SnowDetectionAgent is responsible for analyzing images to determine the presence of snow.
    The agent also generates a unique catalog_guid for each image, which serves as the primary key in the database, based on metadata including creation date, station acronym, location ID, and platform ID.

    Usage:
    - To generate a unique identifier (`catalog_guid`) for each image, use the `generate_unique_id` method with parameters like creation date, station acronym, location ID, and platform ID.
    - The generated `catalog_guid` is then used as a primary key in the DuckDB database to store image analysis results, ensuring unique records for each image.

    Attributes:
        catalog_guid (str): Unique identifier for each image, generated using metadata to ensure uniqueness.
        db_path (str): Path to the DuckDB database file where results are stored.
        connection (duckdb.Connection): DuckDB connection to handle database operations.
        model (LogisticRegression): Machine learning model used to fine-tune snow detection.
    """

    def __init__(self, db_path='snow_analysis.duckdb', model_path='snow_model.pkl'):
        self.connection = duckdb.connect(db_path)
        self.connection.execute('''CREATE TABLE IF NOT EXISTS snow_analysis (
            catalog_guid VARCHAR PRIMARY KEY,
            filepath VARCHAR,
            creation_date TIMESTAMP,
            has_snow_presence BOOLEAN,
            regions JSONB,
            false_positive BOOLEAN DEFAULT FALSE
        );''')

        self.model_path = model_path
        if os.path.exists(model_path):
            with open(model_path, 'rb') as model_file:
                self.model = pickle.load(model_file)
        else:
            self.model = LogisticRegression()
            self.train_initial_model()

    def extract_creation_date(self, filepath: str) -> dict:
        creation_date = None
        extension_coordinates = []
        try:
            image = Image.open(filepath)
            exif_data = image._getexif()
            if exif_data is not None:
                width, height = image.size
                extension_coordinates = [[0, 0], [width, 0], [width, height], [0, height]]
                for tag, value in exif_data.items():
                    tag_name = TAGS.get(tag, tag)
                    if tag_name == 'DateTimeOriginal':
                        creation_date = datetime.strptime(value, '%Y:%m:%d %H:%M:%S')
                        return {
                            'year': creation_date.year,
                            'day_of_year': creation_date.timetuple().tm_yday,
                            'creation_date': creation_date.strftime('%Y-%m-%dT%H:%M'),
                            'extension_coordinates': extension_coordinates
                        }
        except (AttributeError, IOError, ValueError) as e:
            print(f"Error reading EXIF data: {e}")
        except Exception as e:
            print(f"Unexpected error: {e}")
        
        filename = os.path.basename(filepath)
        patterns = [
            r'(\d{4}-\d{2}-\d{2})T(\d{4})\.jpg$',
            r'_(\d{4}-\d{2}-\d{2})_(\d{4})\.jpg$',
            r'_(\d{6})(\d{4})\.jpg$',
            r'_(\d{4}-\d{2}-\d{2})T(\d{2}):(\d{2}):(\d{2})_.*\.jpg$'
        ]
        
        for pattern in patterns:
            match = re.search(pattern, filename)
            if match:
                try:
                    if len(match.groups()) == 2:
                        date_str, time_str = match.groups()
                        if '-' in date_str:
                            creation_date = datetime.strptime(date_str + time_str, '%Y-%m-%d%H%M')
                        else:
                            creation_date = datetime.strptime(date_str + time_str, '%y%m%d%H%M')
                    elif len(match.groups()) == 4:
                        date_str, hour, minute, second = match.groups()
                        creation_date = datetime.strptime(f"{date_str} {hour}:{minute}:{second}", '%Y-%m-%d %H:%M:%S')
                    return {
                        'year': creation_date.year,
                        'day_of_year': creation_date.timetuple().tm_yday,
                        'creation_date': creation_date.strftime('%Y-%m-%dT%H:%M'),
                        'extension_coordinates': extension_coordinates
                    }
                except ValueError as ve:
                    print(f"Error parsing filename for date: {ve}")
                    continue
        
        return None

    def generate_unique_id(self, creation_date: str, station_acronym: str, location_id: str, platform_id: str) -> str:
        unique_string = f"{creation_date}_{station_acronym}_{location_id}_{platform_id}"
        unique_id = hashlib.sha256(unique_string.encode()).hexdigest()
        return unique_id

    def list_image_files(self, dirpath: str, extensions: list = ['.jpg', '.jpeg']):
        file_paths = []
        for root, _, files in os.walk(dirpath):
            for file in files:
                if any(file.lower().endswith(ext) for ext in extensions):
                    file_paths.append(os.path.join(root, file))
        return file_paths

    def is_snow_present(self, image, extension_coordinates=None):
        if extension_coordinates:
            x, y, w, h = extension_coordinates
            image = image[y:y+h, x:x+w]

        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        lower_white = np.array([0, 0, 200], dtype=np.uint8)
        upper_white = np.array([180, 40, 255], dtype=np.uint8)

        mask = cv2.inRange(hsv, lower_white, upper_white)
        snow_pixels = np.count_nonzero(mask)
        total_pixels = image.shape[0] * image.shape[1]
        snow_percentage = snow_pixels / total_pixels
        
        return snow_percentage > 0.05, mask

    def is_snow_present_in_roi(self, mask, roi_polygon):
        height, width = mask.shape
        roi_mask = np.zeros((height, width), dtype=np.uint8)
        
        roi_pts = np.array(roi_polygon, np.int32)
        cv2.fillPoly(roi_mask, [roi_pts], 1)
        
        snow_pixels_roi = cv2.bitwise_and(mask, mask, mask=roi_mask)
        roi_snow_count = np.count_nonzero(snow_pixels_roi)
        roi_total_count = np.count_nonzero(roi_mask)
        
        if roi_total_count == 0:
            return False
        
        roi_snow_percentage = roi_snow_count / roi_total_count
        return roi_snow_percentage > 0.05

    def analyze_image(self, filepath, roi_polygons, extension_coordinates=None, station_acronym='', location_id='', platform_id=''):
        creation_date_info = self.extract_creation_date(filepath)
        if creation_date_info is None:
            print(f"Could not extract creation date for {filepath}")
            return None

        creation_date = creation_date_info['creation_date']
        extension_coordinates = creation_date_info['extension_coordinates']
        catalog_guid = self.generate_unique_id(creation_date, station_acronym, location_id, platform_id)

        check_result = self.connection.execute("SELECT * FROM snow_analysis WHERE catalog_guid = ?", (catalog_guid,)).fetchone()
        if check_result is not None:
            return check_result

        image = cv2.imread(filepath)
        if image is None:
            print(f"Could not read image: {filepath}")
            return None

        has_snow_presence, mask = self.is_snow_present(image, extension_coordinates)
        region_results = {}

        for i, roi_polygon in enumerate(roi_polygons, start=1):
            region_key = f"ROI_{i:02}_has_snow_presence"
            region_results[region_key] = self.is_snow_present_in_roi(mask, roi_polygon)

        record = {
            'catalog_guid': catalog_guid,
            'filepath': filepath,
            'creation_date': creation_date,
            'has_snow_presence': has_snow_presence,
            'regions': region_results,
            'false_positive': False
        }

        self.connection.execute(
            "INSERT INTO snow_analysis (catalog_guid, filepath, creation_date, has_snow_presence, regions, false_positive) VALUES (?, ?, ?, ?, ?, ?)",
            (record['catalog_guid'], record['filepath'], record['creation_date'], record['has_snow_presence'], json.dumps(region_results), record['false_positive'])
        )

        return record

    def flag_false_positive(self, filepath, station_acronym='', location_id='', platform_id=''):
        creation_date_info = self.extract_creation_date(filepath)
        if creation_date_info is None:
            print(f"Could not extract creation date for {filepath}")
            return

        creation_date = creation_date_info['creation_date']
        catalog_guid = self.generate_unique_id(creation_date, station_acronym, location_id, platform_id)

        self.connection.execute(
            "UPDATE snow_analysis SET false_positive = TRUE WHERE catalog_guid = ?",
            (catalog_guid,)
        )
        print(f"Image {filepath} has been flagged as a false positive.")

    def adjust_parameters_based_on_feedback(self):
        false_positives = self.connection.execute("SELECT * FROM snow_analysis WHERE false_positive = TRUE").fetchall()
        if not false_positives:
            print("No false positives to adjust parameters for.")
            return

        X = []
        y = []
        for record in false_positives:
            X.append([0, 40, 255])
            y.append(0)

        positive_examples = self.connection.execute("SELECT * FROM snow_analysis WHERE false_positive = FALSE LIMIT 10").fetchall()
        for record in positive_examples:
            X.append([0, 40, 255])
            y.append(1)

        if len(X) > 1:
            self.model.fit(X, y)
            with open(self.model_path, 'wb') as model_file:
                pickle.dump(self.model, model_file)
            print("Model parameters adjusted based on feedback.")
        else:
            print("Not enough data to retrain the model.")

    def train_initial_model(self):
        X = [[0, 40, 255]]
        y = [1]
        self.model.fit(X, y)
        with open(self.model_path, 'wb') as model_file:
            pickle.dump(self.model, model_file)

    def analyze_directory(self, directory, roi_polygons, extension_coordinates=None, station_acronym='', location_id='', platform_id=''):
        image_files = self.list_image_files(directory)
        for image_file in image_files:
            self.analyze_image(image_file, roi_polygons, extension_coordinates, station_acronym, location_id, platform_id)
