"""
Observation Store - Database layer for IOM observations
Supports PostgreSQL with PostGIS for geospatial queries
"""

import json
import os
from datetime import datetime, timedelta
from typing import List, Optional, Dict, Any
from contextlib import contextmanager

# Database imports
import psycopg2
from psycopg2.extras import RealDictCursor
from psycopg2.pool import ThreadedConnectionPool

from observation_models import Observation, ObservationSummary, ObservationTrend


class ObservationStore:
    """Store and retrieve observations from database"""
    
    def __init__(self, dsn: Optional[str] = None):
        """
        Initialize observation store
        
        Args:
            dsn: PostgreSQL connection string. If None, uses environment variable.
        """
        if dsn is None:
            dsn = os.environ.get(
                'DATABASE_URL',
                'postgresql://localhost:5432/iom'
            )
        
        self.dsn = dsn
        self.pool = None
        self._init_pool()
        self._init_tables()
    
    def _init_pool(self):
        """Initialize connection pool"""
        self.pool = ThreadedConnectionPool(
            minconn=1,
            maxconn=10,
            dsn=self.dsn
        )
    
    @contextmanager
    def _get_connection(self):
        """Get connection from pool"""
        conn = self.pool.getconn()
        try:
            yield conn
        finally:
            self.pool.putconn(conn)
    
    @contextmanager
    def _get_cursor(self, conn):
        """Get cursor from connection"""
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        try:
            yield cursor
        finally:
            cursor.close()
    
    def _init_tables(self):
        """Initialize database tables"""
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                # Enable PostGIS
                cur.execute("CREATE EXTENSION IF NOT EXISTS postgis;")
                
                # Observations table
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS observations (
                        id VARCHAR(20) PRIMARY KEY,
                        timestamp TIMESTAMPTZ NOT NULL,
                        object_goid VARCHAR(50) NOT NULL,
                        observer VARCHAR(50) NOT NULL,
                        findings JSONB NOT NULL DEFAULT '[]',
                        media JSONB NOT NULL DEFAULT '[]',
                        weather JSONB,
                        ai_analysis JSONB,
                        location GEOGRAPHY(POINT, 4326),
                        created_at TIMESTAMPTZ DEFAULT NOW(),
                        updated_at TIMESTAMPTZ DEFAULT NOW()
                    );
                """)
                
                # Indexes
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_observations_object 
                    ON observations(object_goid);
                """)
                
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_observations_timestamp 
                    ON observations(timestamp DESC);
                """)
                
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_observations_location 
                    ON observations USING GIST(location);
                """)
                
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_observations_findings 
                    ON observations USING GIN(findings);
                """)
                
                # Object state table (latest condition per object)
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS object_states (
                        object_goid VARCHAR(50) PRIMARY KEY,
                        latest_condition INTEGER,
                        latest_observation_id VARCHAR(20),
                        latest_observation_date TIMESTAMPTZ,
                        observation_count INTEGER DEFAULT 0,
                        finding_types JSONB DEFAULT '[]',
                        risk_level DECIMAL(3,1),
                        next_observation_due TIMESTAMPTZ,
                        updated_at TIMESTAMPTZ DEFAULT NOW()
                    );
                """)
                
                conn.commit()
    
    def _generate_id(self) -> str:
        """Generate unique observation ID"""
        year = datetime.now().year
        
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT COUNT(*) as count 
                    FROM observations 
                    WHERE id LIKE %s
                """, (f'OBS-{year}-%',))
                
                result = cur.fetchone()
                seq = (result['count'] + 1) if result else 1
                
                return f"OBS-{year}-{seq:07d}"
    
    def create(self, observation: Observation) -> str:
        """
        Save observation to database
        
        Args:
            observation: Observation object to save
            
        Returns:
            Observation ID
        """
        # Generate ID if not provided
        if not observation.id:
            observation.id = self._generate_id()
        
        # Extract location from first media geotag
        location = None
        if observation.media and observation.media[0].geotag:
            geo = observation.media[0].geotag
            location = f"POINT({geo.lng} {geo.lat})"
        
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                # Insert observation
                cur.execute("""
                    INSERT INTO observations (
                        id, timestamp, object_goid, observer,
                        findings, media, weather, ai_analysis, location
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s::geography)
                """, (
                    observation.id,
                    observation.timestamp,
                    observation.object_goid,
                    observation.observer,
                    json.dumps([f.model_dump() for f in observation.findings]),
                    json.dumps([m.model_dump() for m in observation.media]),
                    json.dumps(observation.weather.model_dump()) if observation.weather else None,
                    json.dumps(observation.ai_analysis.model_dump()) if observation.ai_analysis else None,
                    location
                ))
                
                # Update object state
                self._update_object_state(conn, observation)
                
                conn.commit()
                
                return observation.id
    
    def _update_object_state(self, conn, observation: Observation):
        """Update object state with latest observation"""
        with self._get_cursor(conn) as cur:
            # Extract finding types
            finding_types = list(set(f.type.value for f in observation.findings))
            
            # Extract risk level from AI analysis
            risk_level = None
            if observation.ai_analysis:
                # Map condition to risk (simplified)
                condition = observation.ai_analysis.overall_condition
                risk_level = condition * 2  # 1-5 -> 2-10
            
            # Upsert object state
            cur.execute("""
                INSERT INTO object_states (
                    object_goid, latest_condition, latest_observation_id,
                    latest_observation_date, observation_count, finding_types,
                    risk_level, next_observation_due
                ) VALUES (%s, %s, %s, %s, 1, %s, %s, %s)
                ON CONFLICT (object_goid) DO UPDATE SET
                    latest_condition = EXCLUDED.latest_condition,
                    latest_observation_id = EXCLUDED.latest_observation_id,
                    latest_observation_date = EXCLUDED.latest_observation_date,
                    observation_count = object_states.observation_count + 1,
                    finding_types = EXCLUDED.finding_types,
                    risk_level = EXCLUDED.risk_level,
                    next_observation_due = EXCLUDED.next_observation_due,
                    updated_at = NOW();
            """, (
                observation.object_goid,
                observation.ai_analysis.overall_condition if observation.ai_analysis else None,
                observation.id,
                observation.timestamp,
                json.dumps(finding_types),
                risk_level,
                observation.ai_analysis.next_observation_due if observation.ai_analysis else None
            ))
    
    def get(self, observation_id: str) -> Optional[Observation]:
        """
        Get observation by ID
        
        Args:
            observation_id: Observation ID
            
        Returns:
            Observation object or None
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT * FROM observations WHERE id = %s
                """, (observation_id,))
                
                row = cur.fetchone()
                if not row:
                    return None
                
                return self._row_to_observation(row)
    
    def get_for_object(
        self,
        goid: str,
        limit: int = 100,
        offset: int = 0
    ) -> List[Observation]:
        """
        Get all observations for an object
        
        Args:
            goid: Object GOID
            limit: Maximum number of observations
            offset: Offset for pagination
            
        Returns:
            List of observations
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT * FROM observations 
                    WHERE object_goid = %s
                    ORDER BY timestamp DESC
                    LIMIT %s OFFSET %s
                """, (goid, limit, offset))
                
                rows = cur.fetchall()
                return [self._row_to_observation(row) for row in rows]
    
    def get_summary(self, goid: str) -> Optional[ObservationSummary]:
        """
        Get observation summary for an object
        
        Args:
            goid: Object GOID
            
        Returns:
            Observation summary or None
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT * FROM object_states WHERE object_goid = %s
                """, (goid,))
                
                row = cur.fetchone()
                if not row:
                    return None
                
                return ObservationSummary(
                    object_goid=row['object_goid'],
                    observation_count=row['observation_count'],
                    latest_condition=row['latest_condition'],
                    latest_observation_date=row['latest_observation_date'],
                    finding_types=json.loads(row['finding_types']) if row['finding_types'] else [],
                    risk_trend=self._calculate_trend(goid)
                )
    
    def get_trend(self, goid: str, months: int = 6) -> ObservationTrend:
        """
        Get trend analysis for an object
        
        Args:
            goid: Object GOID
            months: Number of months to analyze
            
        Returns:
            Observation trend
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT 
                        DATE_TRUNC('month', timestamp) as month,
                        AVG((ai_analysis->>'overall_condition')::int) as avg_condition,
                        COUNT(*) as observation_count
                    FROM observations
                    WHERE object_goid = %s
                    AND timestamp >= NOW() - INTERVAL '%s months'
                    GROUP BY DATE_TRUNC('month', timestamp)
                    ORDER BY month
                """, (goid, months))
                
                rows = cur.fetchall()
                
                data_points = [
                    {
                        'month': row['month'].strftime('%Y-%m'),
                        'condition': round(row['avg_condition'], 1) if row['avg_condition'] else None,
                        'observations': row['observation_count']
                    }
                    for row in rows
                ]
                
                return ObservationTrend(
                    object_goid=goid,
                    period_months=months,
                    data_points=data_points
                )
    
    def find_by_location(
        self,
        lat: float,
        lng: float,
        radius_meters: float,
        limit: int = 100
    ) -> List[Observation]:
        """
        Find observations near a location
        
        Args:
            lat: Latitude
            lng: Longitude
            radius_meters: Search radius in meters
            limit: Maximum results
            
        Returns:
            List of observations
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT * FROM observations
                    WHERE ST_DWithin(
                        location::geography,
                        ST_SetSRID(ST_MakePoint(%s, %s), 4326)::geography,
                        %s
                    )
                    ORDER BY timestamp DESC
                    LIMIT %s
                """, (lng, lat, radius_meters, limit))
                
                rows = cur.fetchall()
                return [self._row_to_observation(row) for row in rows]
    
    def find_by_defect(
        self,
        defect_code: str,
        min_confidence: float = 0.5,
        limit: int = 100
    ) -> List[Observation]:
        """
        Find observations with specific defect
        
        Args:
            defect_code: Defect code
            min_confidence: Minimum confidence threshold
            limit: Maximum results
            
        Returns:
            List of observations
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    SELECT * FROM observations
                    WHERE findings @> '[{"code": "%s"}]'::jsonb
                    AND EXISTS (
                        SELECT 1 FROM jsonb_array_elements(findings) as f
                        WHERE (f->>'code') = %s
                        AND (f->>'confidence')::float >= %s
                    )
                    ORDER BY timestamp DESC
                    LIMIT %s
                """, (defect_code, defect_code, min_confidence, limit))
                
                rows = cur.fetchall()
                return [self._row_to_observation(row) for row in rows]
    
    def _row_to_observation(self, row: Dict) -> Observation:
        """Convert database row to Observation object"""
        from observation_models import Finding, Media, GeoTag, AIAnalysis, Weather
        
        # Parse findings
        findings_data = json.loads(row['findings']) if row['findings'] else []
        findings = [Finding(**f) for f in findings_data]
        
        # Parse media
        media_data = json.loads(row['media']) if row['media'] else []
        media = []
        for m in media_data:
            if m.get('geotag'):
                m['geotag'] = GeoTag(**m['geotag'])
            media.append(Media(**m))
        
        # Parse weather
        weather = None
        if row['weather']:
            weather = Weather(**json.loads(row['weather']))
        
        # Parse AI analysis
        ai_analysis = None
        if row['ai_analysis']:
            ai_analysis = AIAnalysis(**json.loads(row['ai_analysis']))
        
        return Observation(
            id=row['id'],
            timestamp=row['timestamp'],
            object_goid=row['object_goid'],
            observer=row['observer'],
            findings=findings,
            media=media,
            weather=weather,
            ai_analysis=ai_analysis
        )
    
    def _calculate_trend(self, goid: str) -> Optional[str]:
        """Calculate trend direction for an object"""
        trend = self.get_trend(goid, months=3)
        
        if len(trend.data_points) < 2:
            return None
        
        conditions = [dp['condition'] for dp in trend.data_points if dp['condition'] is not None]
        
        if len(conditions) < 2:
            return None
        
        first = conditions[0]
        last = conditions[-1]
        
        if last < first:
            return "improving"
        elif last > first:
            return "degrading"
        else:
            return "stable"
    
    def get_degrading_objects(
        self,
        threshold: float = 1.0,
        limit: int = 100
    ) -> List[Dict]:
        """
        Find objects that are degrading over time
        
        Args:
            threshold: Minimum condition change to flag
            limit: Maximum results
            
        Returns:
            List of degrading objects with details
        """
        with self._get_connection() as conn:
            with self._get_cursor(conn) as cur:
                cur.execute("""
                    WITH condition_changes AS (
                        SELECT 
                            object_goid,
                            MIN(timestamp) as first_obs,
                            MAX(timestamp) as last_obs,
                            MIN((ai_analysis->>'overall_condition')::int) as min_condition,
                            MAX((ai_analysis->>'overall_condition')::int) as max_condition,
                            COUNT(*) as obs_count
                        FROM observations
                        WHERE timestamp >= NOW() - INTERVAL '6 months'
                        GROUP BY object_goid
                        HAVING COUNT(*) >= 2
                    )
                    SELECT 
                        object_goid,
                        first_obs,
                        last_obs,
                        min_condition as first_condition,
                        max_condition as last_condition,
                        (max_condition - min_condition) as condition_change,
                        obs_count
                    FROM condition_changes
                    WHERE (max_condition - min_condition) >= %s
                    ORDER BY condition_change DESC
                    LIMIT %s
                """, (threshold, limit))
                
                return [dict(row) for row in cur.fetchall()]


if __name__ == '__main__':
    # Example usage
    store = ObservationStore()
    
    from observation_models import Observation, Finding, Media, GeoTag, AIAnalysis
    from datetime import datetime
    
    # Create sample observation
    obs = Observation(
        object_goid="BYG-FAC-WIN-GLA-0001",
        observer="zoomer:anna_k",
        findings=[
            Finding(
                type="dirt_accumulation",
                code="2100",
                description="Smuts på fönster",
                measurement="30% coverage",
                confidence=0.94
            )
        ],
        media=[
            Media(
                type="image",
                url="https://storage.quixzoom.com/obs/img_0001.jpg",
                geotag=GeoTag(lat=59.3293, lng=18.0686, accuracy=2.1)
            )
        ],
        ai_analysis=AIAnalysis(
            model="infrastructure-v3.2",
            overall_condition=3,
            recommended_action="schedule_cleaning"
        )
    )
    
    # Save
    obs_id = store.create(obs)
    print(f"Created observation: {obs_id}")
    
    # Retrieve
    retrieved = store.get(obs_id)
    print(f"Retrieved: {retrieved.object_goid if retrieved else 'Not found'}")